diff --git a/rosie.py b/rosie.py index d39dad4..9ddc5de 100755 --- a/rosie.py +++ b/rosie.py @@ -10,10 +10,9 @@ def entered_command(argv): def help(): message = ( 'Usage:', - ' python rosie.py run chamber_of_deputies []', + ' python rosie.py run (chamber_of_deputies | federal_senate) [] [--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(): diff --git a/rosie/chamber_of_deputies/__init__.py b/rosie/chamber_of_deputies/__init__.py index 9a8989e..367034e 100644 --- a/rosie/chamber_of_deputies/__init__.py +++ b/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() diff --git a/rosie/chamber_of_deputies/adapter.py b/rosie/chamber_of_deputies/adapter.py index b996347..983cd96 100644 --- a/rosie/chamber_of_deputies/adapter.py +++ b/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() diff --git a/rosie/federal_senate/__init__.py b/rosie/federal_senate/__init__.py index 29489d8..5270b8a 100644 --- a/rosie/federal_senate/__init__.py +++ b/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() \ No newline at end of file + core() diff --git a/rosie/federal_senate/adapter.py b/rosie/federal_senate/adapter.py index 22c0c04..1cfdfbb 100644 --- a/rosie/federal_senate/adapter.py +++ b/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()