Linting datasets module & upgrading ruff (#351)

* linting datasets module & upgrading ruff

Author: anupkalburgi

dbldatagen/datasets/__init__.py

@@ -1,20 +1,32 @@
-from .dataset_provider import DatasetProvider, dataset_definition
 from .basic_geometries import BasicGeometriesProvider
 from .basic_process_historian import BasicProcessHistorianProvider
 from .basic_stock_ticker import BasicStockTickerProvider
 from .basic_telematics import BasicTelematicsProvider
 from .basic_user import BasicUserProvider
 from .benchmark_groupby import BenchmarkGroupByProvider
+from .dataset_provider import DatasetProvider, dataset_definition
 from .multi_table_sales_order_provider import MultiTableSalesOrderProvider
 from .multi_table_telephony_provider import MultiTableTelephonyProvider
 
-__all__ = ["dataset_provider",
+
+__all__ = [
+           "BasicGeometriesProvider",
+           "BasicProcessHistorianProvider",
+           "BasicStockTickerProvider",
+           "BasicTelematicsProvider",
+           "BasicUserProvider",
+           "BenchmarkGroupByProvider",
+           "DatasetProvider",
+           "MultiTableSalesOrderProvider",
+           "MultiTableTelephonyProvider",
            "basic_geometries",
            "basic_process_historian",
            "basic_stock_ticker",
            "basic_telematics",
            "basic_user",
            "benchmark_groupby",
+           "dataset_definition",
+           "dataset_provider",
            "multi_table_sales_order_provider",
            "multi_table_telephony_provider"
            ]

dbldatagen/datasets/basic_geometries.py

@@ -1,4 +1,11 @@
-from .dataset_provider import DatasetProvider, dataset_definition
+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 dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/geometries",
@@ -34,7 +41,7 @@ class BasicGeometriesProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
     DEFAULT_MIN_LON = -180.0
     DEFAULT_MAX_LON = 180.0
     COLUMN_COUNT = 2
-    ALLOWED_OPTIONS = [
+    ALLOWED_OPTIONS: ClassVar[list[str]]  = [
         "geometryType",
         "maxVertices",
         "minLatitude",
@@ -45,11 +52,7 @@ class BasicGeometriesProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions=-1,
-                          **options):
-        import dbldatagen as dg
-        import warnings as w
-
+    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)
@@ -72,7 +75,7 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
         )
         if geometryType == "point":
             if maxVertices > 1:
-                w.warn('Ignoring property maxVertices for point geometries')
+                w.warn("Ignoring property maxVertices for point geometries", stacklevel=2)
             df_spec = (
                 df_spec.withColumn("lat", "float", minValue=minLatitude, maxValue=maxLatitude,
                                    step=1e-5, random=generateRandom, omit=True)
@@ -83,7 +86,7 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
         elif geometryType == "lineString":
             if maxVertices < 2:
                 maxVertices = 2
-                w.warn("Parameter maxVertices must be >=2 for 'lineString' geometries; Setting to 2")
+                w.warn("Parameter maxVertices must be >=2 for 'lineString' geometries; Setting to 2", stacklevel=2)
             j = 0
             while j < maxVertices:
                 df_spec = (
@@ -101,7 +104,7 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
         elif geometryType == "polygon":
             if maxVertices < 3:
                 maxVertices = 3
-                w.warn("Parameter maxVertices must be >=3 for 'polygon' geometries; Setting to 3")
+                w.warn("Parameter maxVertices must be >=3 for 'polygon' geometries; Setting to 3", stacklevel=2)
             j = 0
             while j < maxVertices:
                 df_spec = (
@@ -111,7 +114,7 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
                                         step=1e-5, random=generateRandom, omit=True)
                 )
                 j = j + 1
-            vertexIndices = list(range(maxVertices)) + [0]
+            vertexIndices = [*list(range(maxVertices)), 0]
             concatCoordinatesExpr = [f"concat(lon_{j}, ' ', lat_{j}, ', ')" for j in vertexIndices]
             concatPairsExpr = f"replace(concat('POLYGON(', {', '.join(concatCoordinatesExpr)}, ')'), ', )', ')')"
             df_spec = (

dbldatagen/datasets/basic_process_historian.py

@@ -1,4 +1,11 @@
-from .dataset_provider import DatasetProvider, dataset_definition
+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 dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/process_historian",
@@ -43,28 +50,27 @@ class BasicProcessHistorianProvider(DatasetProvider.NoAssociatedDatasetsMixin, D
     DEFAULT_START_TIMESTAMP = "2024-01-01 00:00:00"
     DEFAULT_END_TIMESTAMP = "2024-02-01 00:00:00"
     COLUMN_COUNT = 10
-    ALLOWED_OPTIONS = [
+    ALLOWED_OPTIONS: ClassVar[list[str]] = [
         "numDevices",
         "numPlants",
-        "numTags", 
-        "startTimestamp", 
-        "endTimestamp", 
+        "numTags",
+        "startTimestamp",
+        "endTimestamp",
         "dataQualityRatios",
         "random"
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions=-1, **options):
-        import dbldatagen as dg  # import locally to avoid circular imports
-        import numpy as np
+    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)
         numPlants = options.get("numPlants", self.DEFAULT_NUM_PLANTS)
         numTags = options.get("numTags", self.DEFAULT_NUM_TAGS)
         startTimestamp = options.get("startTimestamp", self.DEFAULT_START_TIMESTAMP)
         endTimestamp = options.get("endTimestamp", self.DEFAULT_END_TIMESTAMP)
-        dataQualityRatios = options.get("dataQualityRatios", None)
+        dataQualityRatios = options.get("dataQualityRatios")
 
         assert tableName is None or tableName == DatasetProvider.DEFAULT_TABLE_NAME, "Invalid table name"
         if rows is None or rows < 0:
@@ -83,7 +89,7 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
             .withColumn("device_id", "string", format="0x%09x", baseColumn="internal_device_id")
             .withColumn("plant_id", "string", values=plant_ids, baseColumn="internal_device_id")
             .withColumn("tag_name", "string", values=tag_names, baseColumn="internal_device_id")
-            .withColumn("ts", "timestamp", begin=startTimestamp, end=endTimestamp, 
+            .withColumn("ts", "timestamp", begin=startTimestamp, end=endTimestamp,
                             interval="1 second", random=generateRandom)
             .withColumn("value", "float", minValue=self.MIN_PROPERTY_VALUE, maxValue=self.MAX_PROPERTY_VALUE,
                              step=1e-3, random=generateRandom)

dbldatagen/datasets/basic_stock_ticker.py

@@ -1,6 +1,11 @@
 from random import random
+from typing import ClassVar
 
-from .dataset_provider import DatasetProvider, dataset_definition
+from pyspark.sql import SparkSession
+
+import dbldatagen as dg
+from dbldatagen.data_generator import DataGenerator
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/stock_ticker",
@@ -21,7 +26,6 @@ class BasicStockTickerProvider(DatasetProvider.NoAssociatedDatasetsMixin, Datase
         - numSymbols: number of unique stock ticker symbols
         - startDate: first date for stock ticker data
         - endDate: last date for stock ticker data
-        
     As the data specification is a DataGenerator object, you can add further columns to the data set and
     add constraints (when the feature is available)
 
@@ -32,14 +36,13 @@ class BasicStockTickerProvider(DatasetProvider.NoAssociatedDatasetsMixin, Datase
     DEFAULT_NUM_SYMBOLS = 100
     DEFAULT_START_DATE = "2024-10-01"
     COLUMN_COUNT = 8
-    ALLOWED_OPTIONS = [
+    ALLOWED_OPTIONS: ClassVar[list[str]] = [
         "numSymbols",
         "startDate"
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions=-1, **options):
-        import dbldatagen as dg
+    def getTableGenerator(self, sparkSession: SparkSession, *, tableName: str|None=None, rows: int=-1, partitions: int=-1, **options: object) -> DataGenerator:
 
         numSymbols = options.get("numSymbols", self.DEFAULT_NUM_SYMBOLS)
         startDate = options.get("startDate", self.DEFAULT_START_DATE)
@@ -59,7 +62,7 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
             .withColumn("rand_value", "float", minValue=0.0, maxValue=1.0, step=0.1,
                         baseColumn="symbol_id", omit=True)
             .withColumn("symbol", "string",
-                        expr="""concat_ws('', transform(split(conv(symbol_id, 10, 26), ''), 
+                        expr="""concat_ws('', transform(split(conv(symbol_id, 10, 26), ''),
                             x -> case when ascii(x) < 10 then char(ascii(x) - 48 + 65) else char(ascii(x) + 10) end))""")
             .withColumn("days_from_start_date", "int", expr=f"floor(try_divide(id, {numSymbols}))", omit=True)
             .withColumn("post_date", "date", expr=f"date_add(cast('{startDate}' as date), days_from_start_date)")
@@ -76,13 +79,13 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
                         expr="case when sin(id % 17) > 0 then -1.0 else 1.0 end",
                         omit=True)
             .withColumn("open_base", "decimal(11,2)",
-                        expr=f"""start_value 
-                            + (volatility * prev_modifier_sign * start_value * sin((id - {numSymbols}) % 17)) 
+                        expr=f"""start_value
+                            + (volatility * prev_modifier_sign * start_value * sin((id - {numSymbols}) % 17))
                             + (growth_rate * start_value * try_divide(days_from_start_date - 1, 365))""",
                         omit=True)
             .withColumn("close_base", "decimal(11,2)",
-                        expr="""start_value 
-                            + (volatility * start_value * sin(id % 17)) 
+                        expr="""start_value
+                            + (volatility * start_value * sin(id % 17))
                             + (growth_rate * start_value * try_divide(days_from_start_date, 365))""",
                         omit=True)
             .withColumn("high_base", "decimal(11,2)",

dbldatagen/datasets/basic_telematics.py

@@ -1,4 +1,11 @@
-from .dataset_provider import DatasetProvider, dataset_definition
+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 dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/telematics",
@@ -24,7 +31,7 @@ class BasicTelematicsProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
         - minLon: minimum longitude
         - maxLon: maximum longitude
         - generateWKT: if `True`, generates the well-known text representation of the location
-        
+
     As the data specification is a DataGenerator object, you can add further columns to the data set and
     add constraints (when the feature is available)
 
@@ -42,7 +49,7 @@ class BasicTelematicsProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
     DEFAULT_MIN_LON = -180.0
     DEFAULT_MAX_LON = 180.0
     COLUMN_COUNT = 6
-    ALLOWED_OPTIONS = [
+    ALLOWED_OPTIONS: ClassVar[list[str]] = [
         "numDevices",
         "startTimestamp",
         "endTimestamp",
@@ -55,10 +62,7 @@ class BasicTelematicsProvider(DatasetProvider.NoAssociatedDatasetsMixin, Dataset
     ]
 
     @DatasetProvider.allowed_options(options=ALLOWED_OPTIONS)
-    def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions=-1,
-                 **options):
-        import dbldatagen as dg
-        import warnings as w
+    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)
@@ -77,52 +81,52 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
             partitions = self.autoComputePartitions(rows, self.COLUMN_COUNT)
         if minLat < -90.0:
             minLat = -90.0
-            w.warn("Received an invalid minLat value; Setting to -90.0")
+            w.warn("Received an invalid minLat value; Setting to -90.0", stacklevel=2)
         if minLat > 90.0:
             minLat = 89.0
-            w.warn("Recieved an invalid minLat value; Setting to 89.0")
+            w.warn("Recieved an invalid minLat value; Setting to 89.0", stacklevel=2)
         if maxLat < -90:
             maxLat = -89.0
-            w.warn("Recieved an invalid maxLat value; Setting to -89.0")
+            w.warn("Recieved an invalid maxLat value; Setting to -89.0", stacklevel=2)
         if maxLat > 90.0:
             maxLat = 90.0
-            w.warn("Received an invalid maxLat value; Setting to 90.0")
+            w.warn("Received an invalid maxLat value; Setting to 90.0", stacklevel=2)
         if minLon < -180.0:
             minLon = -180.0
-            w.warn("Received an invalid minLon value; Setting to -180.0")
+            w.warn("Received an invalid minLon value; Setting to -180.0", stacklevel=2)
         if minLon > 180.0:
             minLon = 179.0
-            w.warn("Received an invalid minLon value; Setting to 179.0")
+            w.warn("Received an invalid minLon value; Setting to 179.0", stacklevel=2)
         if maxLon < -180.0:
             maxLon = -179.0
-            w.warn("Received an invalid maxLon value; Setting to -179.0")
+            w.warn("Received an invalid maxLon value; Setting to -179.0", stacklevel=2)
         if maxLon > 180.0:
             maxLon = 180.0
-            w.warn("Received an invalid maxLon value; Setting to 180.0")
+            w.warn("Received an invalid maxLon value; Setting to 180.0", stacklevel=2)
         if minLon > maxLon:
             (minLon, maxLon) = (maxLon, minLon)
-            w.warn("Received minLon > maxLon; Swapping values")
+            w.warn("Received minLon > maxLon; Swapping values", stacklevel=2)
         if minLat > maxLat:
             (minLat, maxLat) = (maxLat, minLat)
-            w.warn("Received minLat > maxLat; Swapping values")
+            w.warn("Received minLat > maxLat; Swapping values", stacklevel=2)
         df_spec = (
              dg.DataGenerator(sparkSession=sparkSession, rows=rows,
                               partitions=partitions, randomSeedMethod="hash_fieldname")
-            .withColumn("device_id", "long", minValue=self.MIN_DEVICE_ID, maxValue=self.MAX_DEVICE_ID, 
+            .withColumn("device_id", "long", minValue=self.MIN_DEVICE_ID, maxValue=self.MAX_DEVICE_ID,
                             uniqueValues=numDevices, random=generateRandom)
-            .withColumn("ts", "timestamp", begin=startTimestamp, end=endTimestamp, 
+            .withColumn("ts", "timestamp", begin=startTimestamp, end=endTimestamp,
                             interval="1 second", random=generateRandom)
             .withColumn("base_lat", "float", minValue=minLat, maxValue=maxLat, step=0.5,
-                            baseColumn='device_id', omit=True)
+                            baseColumn="device_id", omit=True)
             .withColumn("base_lon", "float", minValue=minLon, maxValue=maxLon, step=0.5,
-                            baseColumn='device_id', omit=True)
+                            baseColumn="device_id", omit=True)
             .withColumn("unv_lat", "float", expr="base_lat + (0.5-format_number(rand(), 3))*1e-3", omit=True)
             .withColumn("unv_lon", "float", expr="base_lon + (0.5-format_number(rand(), 3))*1e-3", omit=True)
-            .withColumn("lat", "float", expr=f"""CASE WHEN unv_lat > {maxLat} THEN {maxLat} 
-                ELSE CASE WHEN unv_lat < {minLat} THEN {minLat} 
+            .withColumn("lat", "float", expr=f"""CASE WHEN unv_lat > {maxLat} THEN {maxLat}
+                ELSE CASE WHEN unv_lat < {minLat} THEN {minLat}
                 ELSE unv_lat END END""")
-            .withColumn("lon", "float", expr=f"""CASE WHEN unv_lon > {maxLon} THEN {maxLon} 
-                ELSE CASE WHEN unv_lon < {minLon} THEN {minLon} 
+            .withColumn("lon", "float", expr=f"""CASE WHEN unv_lon > {maxLon} THEN {maxLon}
+                ELSE CASE WHEN unv_lon < {minLon} THEN {minLon}
                 ELSE unv_lon END END""")
             .withColumn("heading", "integer", minValue=0, maxValue=359, step=1, random=generateRandom)
             .withColumn("wkt", "string", expr="concat('POINT(', lon, ' ', lat, ')')", omit=not generateWkt)

dbldatagen/datasets/basic_user.py

@@ -1,4 +1,10 @@
-from .dataset_provider import DatasetProvider, dataset_definition
+from typing import Any
+
+from pyspark.sql import SparkSession
+
+import dbldatagen as dg
+from dbldatagen.data_generator import DataGenerator
+from dbldatagen.datasets.dataset_provider import DatasetProvider, dataset_definition
 
 
 @dataset_definition(name="basic/user", summary="Basic User Data Set", autoRegister=True, supportsStreaming=True)
@@ -27,10 +33,7 @@ class BasicUserProvider(DatasetProvider.NoAssociatedDatasetsMixin, DatasetProvid
     COLUMN_COUNT = 5
 
     @DatasetProvider.allowed_options(options=["random", "dummyValues"])
-    def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions=-1,
-                          **options):
-        import dbldatagen as dg
-
+    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)
 
@@ -47,13 +50,13 @@ def getTableGenerator(self, sparkSession, *, tableName=None, rows=-1, partitions
                              randomSeedMethod="hash_fieldname")
             .withColumn("customer_id", "long", minValue=1000000, maxValue=self.MAX_LONG, random=generateRandom)
             .withColumn("name", "string",
-                        template=r'\w \w|\w \w \w', random=generateRandom)
+                        template=r"\w \w|\w \w \w", random=generateRandom)
             .withColumn("email", "string",
-                        template=r'\w.\w@\w.com|\w@\w.co.u\k', random=generateRandom)
+                        template=r"\w.\w@\w.com|\w@\w.co.u\k", random=generateRandom)
             .withColumn("ip_addr", "string",
-                        template=r'\n.\n.\n.\n', random=generateRandom)
+                        template=r"\n.\n.\n.\n", random=generateRandom)
             .withColumn("phone", "string",
-                        template=r'(ddd)-ddd-dddd|1(ddd) ddd-dddd|ddd ddddddd',
+                        template=r"(ddd)-ddd-dddd|1(ddd) ddd-dddd|ddd ddddddd",
                         random=generateRandom)
         )