Techdebt: Refactoring pipelines to inject dependencies

stocklake/nasdaqapi/pipeline.py

@@ -1,11 +1,10 @@
 import logging
 import os
-from typing import Any, Optional
+from typing import Optional
 
-# from stocklake.core.base_data_loader import BaseDataLoader
+from stocklake.core.base_data_loader import BaseDataLoader
 from stocklake.core.base_pipeline import BasePipeline
-
-# from stocklake.core.base_preprocessor import BasePreprocessor
+from stocklake.core.base_preprocessor import BasePreprocessor
 from stocklake.core.constants import DATA_DIR
 from stocklake.core.stdout import PrettyStdoutPrint
 from stocklake.exceptions import StockLoaderException
@@ -52,40 +51,44 @@ def save_dir_path(self) -> str:
 
     def run(self):
         logger.info("{} NASDAQ pipeline starts {}".format("=" * 30, "=" * 30))
+        repository = LocalArtifactRepository(self._save_dir)
+
         if self.exchange == Exchange.NASDAQ or self.exchange is None:
             self.stdout.step_start(f"{Exchange.NASDAQ} symbols with nasdapapi")
-            exchange_repo = LocalArtifactRepository(
-                os.path.join(self._save_dir, Exchange.NASDAQ)
+            data_loader = NASDAQSymbolsDataLoader(repository)
+            preprocessor = NASDAQSymbolsPreprocessor(
+                repository, data_loader.artifact_path
             )
-            self._run(exchange_repo, NASDAQSymbolsDataLoader, NASDAQSymbolsPreprocessor)
+            self._run(repository, data_loader, preprocessor)
 
         if self.exchange == Exchange.NYSE or self.exchange is None:
             self.stdout.step_start(f"{Exchange.NYSE} symbols with nasdapapi")
-            exchange_repo = LocalArtifactRepository(
-                os.path.join(self._save_dir, Exchange.NASDAQ)
+            data_loader = NYSESymbolsDataLoader(repository)
+            preprocessor = NYSESymbolsPreprocessor(
+                repository, data_loader.artifact_path
             )
-            self._run(exchange_repo, NYSESymbolsDataLoader, NYSESymbolsPreprocessor)
+            self._run(repository, data_loader, preprocessor)
 
         if self.exchange == Exchange.AMEX or self.exchange is None:
             self.stdout.step_start(f"{Exchange.AMEX} symbols with nasdapapi")
-            exchange_repo = LocalArtifactRepository(
-                os.path.join(self._save_dir, Exchange.NASDAQ)
+            data_loader = AMEXSymbolsDataLoader(repository)
+            preprocessor = AMEXSymbolsPreprocessor(
+                repository, data_loader.artifact_path
             )
-            self._run(exchange_repo, AMEXSymbolsDataLoader, AMEXSymbolsPreprocessor)
+            self._run(repository, data_loader, preprocessor)
 
     def _run(
         self,
         repository: ArtifactRepository,
-        data_loader: Any,
-        preprocessor: Any,
+        data_loader: BaseDataLoader,
+        preprocessor: BasePreprocessor,
     ):
-        _data_loader = data_loader(repository)
         if not self.skip_download:
             self.stdout.normal_message("- Downloading ...")
-            _data_loader.download()
+            data_loader.download()
         else:
             self.stdout.warning_message("- Skip Downloading")
-        preprocessor(repository, _data_loader.artifact_path).process()
+        preprocessor.process()
         self.stdout.success_message(
-            f"- Completed🐳. The artifact is saved to {preprocessor.artifact_path}"
+            f"- Completed🐳. The artifact is saved to {self._save_dir}"
         )

tests/nasdaqapi/test_pipeline.py

@@ -29,9 +29,12 @@ def test_run_each_symbols_with_local_artifact(mock_get, exchange_name, tmpdir):
         data_dir=tmpdir,
     )
     pipeline.run()
-    data_dir = os.path.join(pipeline.save_dir_path, "nasdaq")
-    assert os.path.exists(os.path.join(data_dir, f"{exchange_name}_data.csv"))
-    assert os.path.exists(os.path.join(data_dir, f"raw_{exchange_name}_data.json"))
+    assert os.path.exists(
+        os.path.join(pipeline.save_dir_path, f"{exchange_name}_data.csv")
+    )
+    assert os.path.exists(
+        os.path.join(pipeline.save_dir_path, f"raw_{exchange_name}_data.json")
+    )
 
 
 @mock.patch("requests.get", side_effect=mock_requests_get)
@@ -42,7 +45,10 @@ def test_run_with_local_artifact(mock_get, tmpdir):
         data_dir=tmpdir,
     )
     pipeline.run()
-    data_dir = os.path.join(pipeline.save_dir_path, "nasdaq")
     for exchange_name in Exchange.exchanges():
-        assert os.path.exists(os.path.join(data_dir, f"{exchange_name}_data.csv"))
-        assert os.path.exists(os.path.join(data_dir, f"raw_{exchange_name}_data.json"))
+        assert os.path.exists(
+            os.path.join(pipeline.save_dir_path, f"{exchange_name}_data.csv")
+        )
+        assert os.path.exists(
+            os.path.join(pipeline.save_dir_path, f"raw_{exchange_name}_data.json")
+        )