removing ann101/2 as they are redundant, remove some ignores, remove relative imports

Author: anupkalburgi

dbldatagen/datasets/basic_geometries.py

@@ -1,10 +1,11 @@
-from typing import ClassVar
+import warnings as w
+from typing import Any, ClassVar
 
 from pyspark.sql import SparkSession
 
+import dbldatagen as dg
 from dbldatagen.data_generator import DataGenerator
-
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/geometries",
@@ -51,11 +52,7 @@ class BasicGeometriesProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg # noqa: PLC0415
-        import warnings as w # noqa: PLC0415
-
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
         generateRandom = options.get("random", False)
         geometryType = options.get("geometryType", "point")
         maxVertices = options.get("maxVertices", 1 if geometryType == "point" else 3)

dbldatagen/datasets/basic_process_historian.py

@@ -1,10 +1,11 @@
-from typing import ClassVar
+from typing import Any, ClassVar
 
+import numpy as np
 from pyspark.sql import SparkSession
 
+import dbldatagen as dg
 from dbldatagen.data_generator import DataGenerator
-
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/process_historian",
@@ -60,10 +61,8 @@ class BasicProcessHistorianProvider(DatasetProvider.NoAssociatedDatasetsMixin, D
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg  # noqa: PLC0415  # import locally to avoid circular imports
-        import numpy as np  # noqa: PLC0415
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
+
 
         generateRandom = options.get("random", False)
         numDevices = options.get("numDevices", self.DEFAULT_NUM_DEVICES)

dbldatagen/datasets/basic_stock_ticker.py

@@ -1,11 +1,11 @@
-from typing import ClassVar
 from random import random
+from typing import ClassVar
 
 from pyspark.sql import SparkSession
 
+import dbldatagen as dg
 from dbldatagen.data_generator import DataGenerator
-
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/stock_ticker",
@@ -43,8 +43,6 @@ class BasicStockTickerProvider(DatasetProvider.NoAssociatedDatasetsMixin, Datase
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
     def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg # noqa: PLC0415
 
         numSymbols = options.get("numSymbols", self.DEFAULT_NUM_SYMBOLS)
         startDate = options.get("startDate", self.DEFAULT_START_DATE)

dbldatagen/datasets/basic_telematics.py

@@ -1,9 +1,11 @@
-from typing import ClassVar
+import warnings as w
+from typing import Any, ClassVar
+
 from pyspark.sql import SparkSession
 
+import dbldatagen as dg
 from dbldatagen.data_generator import DataGenerator
-
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/telematics",
@@ -60,11 +62,7 @@ class BasicTelematicsProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        import warnings as w # noqa: PLC0415
-
-        import dbldatagen as dg # noqa: PLC0415
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
 
         generateRandom = options.get("random", False)
         numDevices = options.get("numDevices", self.DEFAULT_NUM_DEVICES)

dbldatagen/datasets/basic_user.py

@@ -1,8 +1,10 @@
+from typing import Any
+
 from pyspark.sql import SparkSession
 
+import dbldatagen as dg
 from dbldatagen.data_generator import DataGenerator
-
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/user", summary="Basic User Data Set", autoRegister=True, supportsStreaming=True)
@@ -31,10 +33,7 @@ class BasicUserProvider(DatasetProvider.NoAssociatedDatasetsMixin, DatasetProvid
     COLUMN_COUNT = 5
 
     @DatasetProvider.allowed_options(options=["random", "dummyValues"])
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg # noqa: PLC0415
-
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
         generateRandom = options.get("random", False)
         dummyValues = options.get("dummyValues", 0)
 

dbldatagen/datasets/benchmark_groupby.py

@@ -1,9 +1,11 @@
-from typing import ClassVar
+import warnings as w
+from typing import Any, ClassVar
+
 from pyspark.sql import SparkSession
 
+import dbldatagen as dg
 from dbldatagen.data_generator import DataGenerator
-
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="benchmark/groupby",
@@ -39,10 +41,7 @@ class BenchmarkGroupByProvider(DatasetProvider.NoAssociatedDatasetsMixin, Datase
     ALLOWED_OPTIONS: ClassVar[list[str]] = ["groups", "percentNulls", "rows", "partitions", "tableName", "random"]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg # noqa: PLC0415
-        import warnings as w # noqa: PLC0415
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
 
         generateRandom = options.get("random", False)
         groups = options.get("groups", self.DEFAULT_NUM_GROUPS)

dbldatagen/datasets/dataset_provider.py

@@ -187,7 +187,7 @@ def getRegisteredDatasetsVersion(cls) -> int   :
         return cls._registeredDatasetsVersion
 
     @abstractmethod
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
         """Gets data generation instance that will produce table for named table
 
         :param sparkSession: Spark session to use
@@ -207,7 +207,7 @@ def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=N
 
     @abstractmethod
     def getAssociatedDataset(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1,
-                             **options: object) -> DataGenerator:
+                             **options: dict[str, Any]) -> DataGenerator:
         """
         Gets associated datasets that are used in conjunction with the provider datasets.
         These may be associated lookup tables, tables that execute benchmarks or exercise key features as part of
@@ -288,7 +288,7 @@ class NoAssociatedDatasetsMixin(ABC):  # noqa: B024
             any associated datasets
         """
         def getAssociatedDataset(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int   =-1,
-                                 **options: object) -> DataGenerator:
+                                 **options: dict[str, Any]) -> DataGenerator:
             raise NotImplementedError("Data provider does not produce any associated datasets!")
 
     class DatasetDecoratorUtils:
@@ -381,7 +381,7 @@ def mkClass(self, autoRegister: bool =False) -> type:
             return retval
 
 
-def dataset_definition(cls: type|None =None, *args: Any, autoRegister: bool =False, **kwargs: Any) -> type:  # pylint: disable=keyword-arg-before-vararg # noqa: ANN401
+def dataset_definition(cls: type|None =None, *args: object, autoRegister: bool =False, **kwargs: object) -> type:
     """ decorator to define standard dataset definition
 
     This is intended to be applied classes derived from DatasetProvider to simplify the implementation
@@ -414,7 +414,7 @@ class X(DatasetProvider)
 
     """
 
-    def inner_wrapper(inner_cls: type|None =None, *inner_args: Any, **inner_kwargs) -> type:  # pylint: disable=keyword-arg-before-vararg # noqa: ANN401
+    def inner_wrapper(inner_cls: type|None =None, *inner_args: object, **inner_kwargs: object) -> type:
         """ The inner wrapper function is used to handle the case where the decorator is used with arguments.
         It defers the application of the decorator to the target class until the target class is available.
 

dbldatagen/datasets/multi_table_sales_order_provider.py

@@ -1,8 +1,11 @@
-from pyspark.sql import SparkSession
+from typing import Any
 
-from dbldatagen.data_generator import DataGenerator
+import pyspark.sql.functions as F
+from pyspark.sql import DataFrame, SparkSession
 
-from .dataset_provider import DatasetProvider, dataset_definition
+import dbldatagen as dg
+from dbldatagen.data_generator import DataGenerator
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="multi_table/sales_order", summary="Multi-table sales order dataset", supportsStreaming=True,
@@ -65,7 +68,7 @@ class MultiTableSalesOrderProvider(DatasetProvider):
     INVOICE_MIN_VALUE = 1_000_000
 
     def getCustomers(self, sparkSession: SparkSession, *, rows: int, partitions: int, numCustomers: int, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
+        import dbldatagen as dg  # noqa: PLC0415
 
         # Validate the options:
         if numCustomers is None or numCustomers < 0:
@@ -106,7 +109,7 @@ def getCustomers(self, sparkSession: SparkSession, *, rows: int, partitions: int
         return customers_data_spec
 
     def getCarriers(self, sparkSession: SparkSession, *, rows: int, partitions: int, numCarriers: int, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
+        import dbldatagen as dg  # noqa: PLC0415
 
         # Validate the options:
         if numCarriers is None or numCarriers < 0:
@@ -143,7 +146,7 @@ def getCarriers(self, sparkSession: SparkSession, *, rows: int, partitions: int,
         return carriers_data_spec
 
     def getCatalogItems(self, sparkSession: SparkSession, *, rows: int, partitions: int, numCatalogItems: int, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
+        import dbldatagen as dg  # noqa: PLC0415
 
         # Validate the options:
         if numCatalogItems is None or numCatalogItems < 0:
@@ -184,7 +187,7 @@ def getCatalogItems(self, sparkSession: SparkSession, *, rows: int, partitions:
 
     def getBaseOrders(self, sparkSession: SparkSession, *, rows: int, partitions: int, numOrders: int, numCustomers: int, startDate: str,
                       endDate: str, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
+        import dbldatagen as dg  # noqa: PLC0415
 
         # Validate the options:
         if numOrders is None or numOrders < 0:
@@ -231,7 +234,7 @@ def getBaseOrders(self, sparkSession: SparkSession, *, rows: int, partitions: in
 
     def getBaseOrderLineItems(self, sparkSession: SparkSession, *, rows: int, partitions: int, numOrders: int, numCatalogItems: int,
                               lineItemsPerOrder: int, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
+        import dbldatagen as dg  # noqa: PLC0415
 
         # Validate the options:
         if numOrders is None or numOrders < 0:
@@ -268,9 +271,6 @@ def getBaseOrderLineItems(self, sparkSession: SparkSession, *, rows: int, partit
         return base_order_line_items_data_spec
 
     def getBaseOrderShipments(self, sparkSession: SparkSession, *, rows: int, partitions: int, numOrders: int, numCarriers: int, dummyValues: int) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg # noqa: PLC0415
-
         # Validate the options:
         if numOrders is None or numOrders < 0:
             numOrders = self.DEFAULT_NUM_ORDERS
@@ -311,10 +311,6 @@ def getBaseOrderShipments(self, sparkSession: SparkSession, *, rows: int, partit
         return base_order_shipments_data_spec
 
     def getBaseInvoices(self, sparkSession: SparkSession, *, rows: int, partitions: int, numOrders: int, dummyValues: int) -> DataGenerator:
-        # ruff: noqa: I001
-        import dbldatagen as dg # noqa: PLC0415
-
-
         # Validate the options:
         if numOrders is None or numOrders < 0:
             numOrders = self.DEFAULT_NUM_ORDERS
@@ -354,7 +350,7 @@ def getBaseInvoices(self, sparkSession: SparkSession, *, rows: int, partitions:
 
     @DatasetProvider.allowed_options(options=["numCustomers", "numCarriers", "numCatalogItems", "numOrders",
                                               "lineItemsPerOrder", "startDate", "endDate", "dummyValues"])
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
         # Get the option values:
         numCustomers = options.get("numCustomers", self.DEFAULT_NUM_CUSTOMERS)
         numCarriers = options.get("numCarriers", self.DEFAULT_NUM_CARRIERS)
@@ -443,11 +439,7 @@ def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=N
         "baseOrderShipments",
         "baseInvoices"
     ])
-    def getAssociatedDataset(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
-        # ruff: noqa: I001
-        from pyspark.sql import DataFrame # noqa: PLC0415
-        import pyspark.sql.functions as F  # noqa: PLC0415
-
+    def getAssociatedDataset(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
         dfCustomers = options.get("customers")
         assert dfCustomers is not None and issubclass(type(dfCustomers), DataFrame), \
             "Option `customers` should be a dataframe of customer records"

dbldatagen/datasets/multi_table_telephony_provider.py

@@ -1,8 +1,10 @@
+from typing import Any
 from pyspark.sql import SparkSession
 
 from dbldatagen.data_generator import DataGenerator
 
-from .dataset_provider import DatasetProvider, dataset_definition
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
+import dbldatagen as dg
 
 
 @dataset_definition(name="multi_table/telephony", summary="Multi-table telephony dataset", supportsStreaming=True,
@@ -60,7 +62,6 @@ class MultiTableTelephonyProvider(DatasetProvider):
     DEFAULT_AVG_EVENTS_PER_CUSTOMER = 50
 
     def getPlans(self, sparkSession: SparkSession, *, rows: int, partitions: int, generateRandom: bool, numPlans: int, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
 
         if numPlans is None or numPlans < 0:
             numPlans = self.DEFAULT_NUM_PLANS
@@ -219,7 +220,7 @@ def getDeviceEvents(self, sparkSession: SparkSession, *, rows: int, partitions:
 
     @DatasetProvider.allowed_options(options=["random", "numPlans", "numCustomers", "dummyValues", "numDays",
                                               "averageEventsPerCustomer"])
-    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: dict[str, Any]) -> DataGenerator:
         generateRandom = options.get("random", False)
         numPlans = options.get("numPlans", self.DEFAULT_NUM_PLANS)
         numCustomers = options.get("numCustomers", self.DEFAULT_NUM_CUSTOMERS)
@@ -240,7 +241,7 @@ def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=N
 
     @DatasetProvider.allowed_options(options=["plans", "customers", "deviceEvents"])
     def getAssociatedDataset(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1,
-                             **options: object) -> DataGenerator:
+                             **options: dict[str, Any]) -> DataGenerator:
         # ruff: noqa: I001
         import pyspark.sql.functions as F #noqa: PLC0415
         from pyspark.sql import DataFrame #noqa: PLC0415
@@ -297,11 +298,7 @@ def getAssociatedDataset(self, sparkSession: SparkSession, *, tableName: str|Non
             """)
 
             df_summary.createOrReplaceTempView("mtp_event_summary")
-
-            df_customer_summary = ( # noqa: F841
-                df_customer_pricing.join(df_summary,
-                                         df_customer_pricing.device_id == df_summary.device_id)
-                .createOrReplaceTempView("mtp_customer_summary"))
+            df_customer_pricing.join(df_summary,df_customer_pricing.device_id == df_summary.device_id).createOrReplaceTempView("mtp_customer_summary")
 
             df_invoices = sparkSession.sql("""
                                  select *,

pyproject.toml

@@ -65,7 +65,7 @@ packages = ["dbldatagen"]
 [tool.hatch.build.targets.sdist]
 include = [
     "/dbldatagen",
-    "/tests", 
+    "/tests",
     "/examples",
     "/tutorial",
     "/docs",
@@ -185,8 +185,6 @@ ignore = [
     "SIM102", # Use a single if-statement
     "SIM108", # Use ternary operator
     "UP007",  # Use X | Y for type annotations (keep Union for compatibility)
-    "ANN101", # Missing type annotation for `self` in method
-    "ANN102", # Missing type annotation for `cls` in method
     "ANN003", # Missing type annotation for **kwargs
 ]