Correcting return type

Removing noqa: PLC0415

Author: anupkalburgi

dbldatagen/datasets/dataset_provider.py

@@ -254,7 +254,7 @@ def wrapper(*args, **kwargs) -> Callable: # noqa: ANN002
 
         return decorator
 
-    def checkOptions(self, options: dict[str, Any], allowedOptions: list[str]) -> None:
+    def checkOptions(self, options: dict[str, Any], allowedOptions: list[str]) -> DatasetDefinition:
         """ Check that options are valid
 
         :param options: options to check as dict

dbldatagen/datasets/multi_table_sales_order_provider.py

@@ -68,8 +68,6 @@ 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
-
         # Validate the options:
         if numCustomers is None or numCustomers < 0:
             numCustomers = self.DEFAULT_NUM_CUSTOMERS
@@ -109,8 +107,6 @@ 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
-
         # Validate the options:
         if numCarriers is None or numCarriers < 0:
             numCarriers = self.DEFAULT_NUM_CARRIERS
@@ -146,9 +142,6 @@ 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
-
-        # Validate the options:
         if numCatalogItems is None or numCatalogItems < 0:
             numCatalogItems = self.DEFAULT_NUM_CATALOG_ITEMS
         if rows is None or rows < 0:
@@ -187,8 +180,6 @@ 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
-
         # Validate the options:
         if numOrders is None or numOrders < 0:
             numOrders = self.DEFAULT_NUM_ORDERS
@@ -234,9 +225,6 @@ 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
-
-        # Validate the options:
         if numOrders is None or numOrders < 0:
             numOrders = self.DEFAULT_NUM_ORDERS
         if numCatalogItems is None or numCatalogItems < 0:

dbldatagen/datasets/multi_table_telephony_provider.py

@@ -1,10 +1,11 @@
 from typing import Any
-from pyspark.sql import SparkSession
 
-from dbldatagen.data_generator import DataGenerator
+import pyspark.sql.functions as F
+from pyspark.sql import DataFrame, SparkSession
 
-from dbldatagen.datasets.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/telephony", summary="Multi-table telephony dataset", supportsStreaming=True,
@@ -107,8 +108,6 @@ def getPlans(self, sparkSession: SparkSession, *, rows: int, partitions: int, ge
         return plan_dataspec
 
     def getCustomers(self, sparkSession: SparkSession, *, rows: int, partitions: int, generateRandom: bool, numCustomers: int, numPlans: int, dummyValues: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
-
         if numCustomers is None or numCustomers < 0:
             numCustomers = self.DEFAULT_NUM_CUSTOMERS
 
@@ -149,7 +148,6 @@ def getCustomers(self, sparkSession: SparkSession, *, rows: int, partitions: int
 
     def getDeviceEvents(self, sparkSession: SparkSession, *, rows: int, partitions: int, generateRandom: bool, numCustomers: int, numDays: int, dummyValues: int,
                         averageEventsPerCustomer: int) -> DataGenerator:
-        import dbldatagen as dg # noqa: PLC0415
         MB_100 = 100 * 1000 * 1000
         K_1 = 1000
 
@@ -242,13 +240,9 @@ 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: dict[str, Any]) -> DataGenerator:
-        # ruff: noqa: I001
-        import pyspark.sql.functions as F #noqa: PLC0415
-        from pyspark.sql import DataFrame #noqa: PLC0415
 
         dfPlans = options.get("plans")
-        assert dfPlans is not None and issubclass(type(dfPlans), DataFrame), \
-            "Option `plans` should be a dataframe of plan records"
+        assert dfPlans is not None and issubclass(type(dfPlans), DataFrame), "Option `plans` should be a dataframe of plan records"
 
         dfCustomers = options.get("customers")
         assert dfCustomers is not None and issubclass(type(dfCustomers), DataFrame), \