Add command line option to limit dataset years

rosie.py

@@ -10,10 +10,9 @@ def entered_command(argv):
 def help():
     message = (
         'Usage:',
-        '  python rosie.py run chamber_of_deputies [<path to output directory>]',
+        '  python rosie.py run (chamber_of_deputies | federal_senate) [<path to output directory>] [--years 2017,2016]',
         'Testing:',
-        '  python rosie.py test',
-        '  python rosie.py test chamber_of_deputies',
+        '  python rosie.py test [chamber_of_deputies | federal_senate]',
     )
     print('\n'.join(message))
 
@@ -30,8 +29,11 @@ def run():
         help()
         exit(1)
     target_directory = argv[3] if len(argv) >= 4 else '/tmp/serenata-data/'
+    if '--years' in argv:
+        years = argv[argv.index('--years') + 1]
+        years = [int(num) for num in years.split(',')]
     klass = getattr(rosie, target_module)
-    klass.main(target_directory)
+    klass.main(target_directory, years=years)
 
 
 def test():

rosie/chamber_of_deputies/__init__.py

@@ -3,7 +3,7 @@
 from rosie.core import Core
 
 
-def main(target_directory='/tmp/serenata-data'):
-    adapter = Adapter(target_directory)
+def main(target_directory='/tmp/serenata-data', years=None):
+    adapter = Adapter(target_directory, years=years)
     core = Core(settings, adapter)
     core()

rosie/chamber_of_deputies/adapter.py

@@ -17,8 +17,9 @@
 class Adapter:
     COMPANIES_DATASET = '2016-09-03-companies.xz'
 
-    def __init__(self, path):
+    def __init__(self, path, years=None):
         self.path = path
+        self.years = years
 
     @property
     def dataset(self):
@@ -56,7 +57,7 @@ def rename_categories(self):
 
     def update_datasets(self):
         os.makedirs(self.path, exist_ok=True)
-        chamber_of_deputies = Dataset(self.path)
+        chamber_of_deputies = Dataset(self.path, self.years)
         chamber_of_deputies.fetch()
         chamber_of_deputies.convert_to_csv()
         chamber_of_deputies.translate()

rosie/federal_senate/__init__.py

@@ -4,6 +4,6 @@
 
 
 def main(target_directory='/tmp/serenata-data'):
-    adapter = Adapter(target_directory)
+    adapter = Adapter(target_directory, years=years)
     core = Core(settings, adapter)
-    core()
+    core()

rosie/federal_senate/adapter.py

@@ -14,8 +14,9 @@
 
 class Adapter:
 
-    def __init__(self, path):
+    def __init__(self, path, years=None):
         self.path = path
+        self.years = years
 
     @property
     def dataset(self):
@@ -43,7 +44,7 @@ def create_columns(self):
 
     def update_datasets(self):
         os.makedirs(self.path, exist_ok=True)
-        federal_senate = Dataset(self.path)
+        federal_senate = Dataset(self.path, years=self.years)
         federal_senate.fetch()
         federal_senate.translate()
         federal_senate_reimbursements_path = federal_senate.clean()