From 26dacda732162e7cadbad66defc33dfc5dc87916 Mon Sep 17 00:00:00 2001
From: Eduardo Cuducos <cuducos@users.noreply.github.com>
Date: Sun, 17 Nov 2019 11:20:24 -0500
Subject: [PATCH] Add companies module to the CLI

---
 README.rst                        | 4 +++-
 serenata_toolbox/serenata-toolbox | 7 +++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index 955b389..9824402 100644
--- a/README.rst
+++ b/README.rst
@@ -48,6 +48,8 @@ But you can specify which datasets to download and where to save them. For examp
 
     $ serenata-toolbox /tmp/serenata-data --module federal_senate chamber_of_deputies
 
+Available modules are ``chamber_of_deputies``, ``companies`` and ``federal_senate``.
+
 Yet, you can specify a specific year:
 
 .. code-block:: bash
@@ -58,7 +60,7 @@ Or use it all together:
 
 .. code-block:: bash
 
-    $ serenata-toolbox /tmp/serenata-data --module federal_senate --year 2017
+    $ serenata-toolbox /tmp/serenata-data --module federal_senate companies --year 2017
 
 Finally, you might want to get help:
 
diff --git a/serenata_toolbox/serenata-toolbox b/serenata_toolbox/serenata-toolbox
index c6b9049..b2d382b 100755
--- a/serenata_toolbox/serenata-toolbox
+++ b/serenata_toolbox/serenata-toolbox
@@ -6,6 +6,7 @@ from argparse import ArgumentParser
 
 from serenata_toolbox import log
 from serenata_toolbox.chamber_of_deputies.reimbursements import Reimbursements as ChamberDataset
+from serenata_toolbox.companies.dataset import Dataset as CompaniesDataset
 from serenata_toolbox.datasets import Datasets
 from serenata_toolbox.federal_senate.dataset import Dataset as SenateDataset
 
@@ -26,6 +27,7 @@ def main():
         datasets = Datasets(path)
         datasets.downloader.download(datasets.downloader.LATEST)
         return
+
     if 'chamber_of_deputies' in args.module:
         if args.year:
             for year in args.year:
@@ -33,6 +35,7 @@ def main():
         else:
             chamber = ChamberDataset(path=path)
         chamber()
+
     if 'federal_senate' in args.module:
         if args.year:
             for year in args.year:
@@ -43,6 +46,10 @@ def main():
         senate.translate()
         senate.clean()
 
+    if 'companies' in args.module:
+        companies = CompaniesDataset(path)
+        companies()
+
 
 if __name__ == '__main__':
     main()