diff --git a/attr_dict.py b/attr_dict.py new file mode 100644 index 0000000..52cad55 --- /dev/null +++ b/attr_dict.py @@ -0,0 +1,53 @@ +#!/usr/bin/python3 + +class AttrDict(dict): + ''' + Allow keys of a dictionary to also be called like class attributes: + + >>> regions = {'Milotic': 'Hoenn'} + >>> regions['Milotic'] + Hoenn + >>> regions.Milotic + AttributeError: 'dict' object has no attribute 'Milotic' + + >>> regions = AttrDict({'Milotic': 'Hoenn'}) + >>> regions['Milotic'] # standard dict behavior + Hoenn + >>> regions.Milotic # behavior enabled by AttrDict + Hoenn + + Also handles setting and deleting keys/attrs. Nested dicts are also + converted to AttrDict instances. Less-than-free-range code. + ''' + def __init__(self, *args, **kwargs): + super().__init__() + self.update(*args, **kwargs) + + # Handle setting new keys, turning any nested dicts into AttrDicts as well + def __setitem__(self, key, val): + if isinstance(val, dict) and not isinstance(val, AttrDict): + val = AttrDict(val) + super().__setitem__(key, val) + + def update(self, *args, **kwargs): + for key, val in dict(*args, **kwargs).items(): + self.__setitem__(key, val) + + # Allow dict keys to be called, set, and deleted like class attributes + def __getattr__(self, name): + try: + return self[name] + except KeyError: + raise AttributeError(name) + + def __setattr__(self, name, val): + # leave handling private variable assignment to parent dict type + if name.startswith('_'): + return super().__setattr__(name, val) + self.__setitem__(name, val) + + def __delattr__(self, name): + try: + del self[name] + except KeyError: + raise AttributeError(name) diff --git a/portfolio_maker.py b/portfolio_maker.py index d8f1bb1..f4ae8e0 100755 --- a/portfolio_maker.py +++ b/portfolio_maker.py @@ -1,4 +1,5 @@ #!/usr/bin/python3 +from attr_dict import AttrDict from datetime import datetime, timedelta from io import BytesIO, TextIOWrapper from zipfile import ZipFile @@ -47,7 +48,8 @@ def __init__(self, sat_frac, relative_core_frac=True): self._valid_tix = self._fetch_ticker_info() # set up initial class attributes - self.assets = {} + self.assets = AttrDict() + # allows keys to be called like attrs (e.g., assets['df'] or assets.df) self.tick_info = pd.DataFrame() # will eventually contain rows of assets selected from valid_tix diff --git a/simulator.py b/simulator.py index 7a180e6..e56840a 100755 --- a/simulator.py +++ b/simulator.py @@ -1,4 +1,5 @@ #!/usr/bin/python3 +from attr_dict import AttrDict from better_abc import ABC, abstract_attribute, abstractmethod from datetime import timedelta from portfolio_maker import PortfolioMaker @@ -120,8 +121,9 @@ def __init__(self, Portfolio, cash=1e4, # track the current simulation date self.today = self.open_date - # validate proposed asset dictionary, then add historical data to it - self.assets = self._validate_assets_dict(Portfolio, verbose) + # validate proposed asset dictionary, then add historical data to it. + # allow keys to be called like attrs (e.g., assets['df'] or assets.df) + self.assets = AttrDict(self._validate_assets_dict(Portfolio, verbose)) # make arrays of all dates in set and all *active* dates self.all_dates, self.active_dates = self._get_date_arrays()