black

ddataflow/data_source.py

@@ -129,6 +129,4 @@ def _estimate_size(self, df: DataFrame) -> float:
 
         print(f"Amount of rows in dataframe to estimate size: {df.count()}")
         average_variable_size_bytes = 50
-        return (df.count() * len(df.columns) * average_variable_size_bytes) / (
-            1024**3
-        )
+        return (df.count() * len(df.columns) * average_variable_size_bytes) / (1024**3)

ddataflow/ddataflow.py

@@ -6,9 +6,13 @@
 from ddataflow.data_sources import DataSources
 from ddataflow.downloader import DataSourceDownloader
 from ddataflow.exceptions import WriterNotFoundException
-from ddataflow.sampling.default import build_default_sampling_for_sources, DefaultSamplerOptions
+from ddataflow.sampling.default import (
+    build_default_sampling_for_sources,
+    DefaultSamplerOptions,
+)
 from ddataflow.sampling.sampler import Sampler
 from ddataflow.utils import get_or_create_spark, using_databricks_connect
+from pyspark.sql import DataFrame
 
 logger = logging.getLogger(__name__)
 handler = logging.StreamHandler()
@@ -72,7 +76,6 @@ def __init__(
         self._snapshot_path = base_path + "/" + project_folder_name
         self._local_path = self._LOCAL_BASE_SNAPSHOT_PATH + "/" + project_folder_name
 
-
         if default_sampler:
             # set this before creating data sources
             DefaultSamplerOptions.set(default_sampler)
@@ -92,7 +95,6 @@ def __init__(
             size_limit=self._size_limit,
         )
 
-
         self._data_writers: dict = data_writers if data_writers else {}
 
         self._offline_enabled = os.getenv(self._ENABLE_OFFLINE_MODE_ENVVARIABLE, False)
@@ -116,15 +118,18 @@ def __init__(
         if self._ddataflow_enabled:
             self.set_logger_level(logging.DEBUG)
         else:
-            logger.info("DDataflow is now DISABLED."
-                        "PRODUCTION data will be used and it will write to production tables.")
+            logger.info(
+                "DDataflow is now DISABLED."
+                "PRODUCTION data will be used and it will write to production tables."
+            )
 
     @staticmethod
     def setup_project():
         """
         Sets up a new ddataflow project with empty data sources in the current directory
         """
         from ddataflow.setup.setup_project import setup_project
+
         setup_project()
 
     @staticmethod
@@ -167,60 +172,61 @@ def current_project() -> "DDataflow":
 
         return ddataflow_config.ddataflow
 
+    def source(self, name: str, debugger=False) -> DataFrame:
+        """
+        Gives access to the data source configured in the dataflow
+
+        You can also use this function in the terminal with --debugger=True to inspect the dataframe.
+        """
+        self.print_status()
+
+        logger.debug("Loading data source")
+        data_source: DataSource = self._data_sources.get_data_source(name)
+        logger.debug("Data source loaded")
+        df = self._get_data_from_data_source(data_source)
+
+        if debugger:
+            logger.debug(f"Debugger enabled: {debugger}")
+            logger.debug("In debug mode now, use query to inspect it")
+            breakpoint()
+
+        return df
+
     def _get_spark(self):
         return get_or_create_spark()
 
     def enable(self):
         """
-        When enabled ddataflow will read from the filtered datasoruces
+        When enabled ddataflow will read from the filtered data sources
         instead of production tables. And write to testing tables instead of production ones.
         """
 
         self._ddataflow_enabled = True
 
-    def is_enabled(self):
+    def is_enabled(self) -> bool:
         return self._ddataflow_enabled
 
     def enable_offline(self):
         """Programatically enable offline mode"""
         self._offline_enabled = True
         self.enable()
 
-    def is_local(self):
+    def is_local(self) -> bool:
         return self._offline_enabled
 
     def disable_offline(self):
         """Programatically enable offline mode"""
         self._offline_enabled = False
 
-    def source(self, name: str, debugger=False):
-        """
-        Gives access to the data source configured in the dataflow
-
-        You can also use this function in the terminal with --debugger=True to inspect the dataframe.
-        """
-        self.print_status()
-
-        logger.debug("Loading data source")
-        data_source: DataSource = self._data_sources.get_data_source(name)
-        logger.debug("Data source loaded")
-        df = self._get_df_from_source(data_source)
-
-        if debugger:
-            logger.debug(f"Debugger enabled: {debugger}")
-            logger.debug("In debug mode now, use query to inspect it")
-            breakpoint()
-
-        return df
-
-    def source_name(self, name, disable_view_creation=False):
+    def source_name(self, name, disable_view_creation=False) -> str:
         """
         Given the name of a production table, returns the name of the corresponding ddataflow table when ddataflow is enabled
         If ddataflow is disabled get the production one.
         """
         logger.debug("Source name used: ", name)
         source_name = name
 
+        # the gist of ddtafalow
         if self._ddataflow_enabled:
             source_name = self._get_new_table_name(name)
             if disable_view_creation:
@@ -286,7 +292,7 @@ def disable(self):
         """Disable ddtaflow overriding tables, uses production state in other words"""
         self._ddataflow_enabled = False
 
-    def _get_df_from_source(self, data_source):
+    def _get_data_from_data_source(self, data_source: DataSource) -> DataFrame:
         if not self._ddataflow_enabled:
             logger.debug("DDataflow not enabled")
             # goes directly to production without prefilters

ddataflow/setup/ddataflow_config.py

@@ -5,25 +5,25 @@
     "project_folder_name": "your_project_name",
     # add here your tables or paths with customized sampling logic
     "data_sources": {
-        #"demo_tours": {
+        # "demo_tours": {
         #    # use default_sampling=True to rely on automatic sampling otherwise it will use the whole data
         #    "default_sampling": True,
-        #},
-        #"demo_tours2": {
+        # },
+        # "demo_tours2": {
         #    "source": lambda spark: spark.table("demo_tours"),
         #    # to customize the sampling logic
         #    "filter": lambda df: df.limit(500),
-        #},
-        #"/mnt/cleaned/EventName": {
+        # },
+        # "/mnt/cleaned/EventName": {
         #    "file-type": "parquet",
         #    "default_sampling": True,
-        #},
+        # },
     },
-    #"default_sampler": {
+    # "default_sampler": {
     #    # defines the amount of rows retrieved with the default sampler, used as .limit(limit) in the dataframe
     #    # default = 10000
     #    "limit": 100000,
-    #},
+    # },
     # to customize the max size of your examples uncomment the line below
     # "data_source_size_limit_gb": 3
     # to customize the location of your test datasets in your data wharehouse

examples/ddataflow_config.py

@@ -4,16 +4,16 @@
     # add here your tables or paths with customized sampling logic
     "data_sources": {
         "demo_tours": {
-            "source": lambda spark: spark.table('demo_tours'),
-            "filter": lambda df: df.limit(500)
+            "source": lambda spark: spark.table("demo_tours"),
+            "filter": lambda df: df.limit(500),
         },
         "demo_locations": {
-            "source": lambda spark: spark.table('demo_locations'),
+            "source": lambda spark: spark.table("demo_locations"),
             "default_sampling": True,
-        }
+        },
     },
     "project_folder_name": "ddataflow_demo",
 }
 
 # initialize the application and validate the configuration
-ddataflow = DDataflow(**config)
+ddataflow = DDataflow(**config)

tests/test_configuration.py

@@ -41,6 +41,7 @@
     "data_source_size_limit_gb": 3,
 }
 
+
 def test_initialize_successfully():
     """
     Tests that a correct config will not fail to be instantiated

tests/test_sql.py

@@ -8,14 +8,8 @@ def test_sql():
 
     config = {
         "project_folder_name": "unit_tests",
-        "data_sources": {
-            "location": {
-                'default_sampling': True
-            }
-        },
-        'default_sampler': {
-            'limit': 2
-        }
+        "data_sources": {"location": {"default_sampling": True}},
+        "default_sampler": {"limit": 2},
     }
     ddataflow = DDataflow(**config)