fixing tests

Author: anupkalburgi

dbldatagen/spec/column_spec.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from .compat import BaseModel, root_validator
+
+
+DbldatagenBasicType = Literal[
+    "string",
+    "int",
+    "long",
+    "float",
+    "double",
+    "decimal",
+    "boolean",
+    "date",
+    "timestamp",
+    "short",
+    "byte",
+    "binary",
+    "integer",
+    "bigint",
+    "tinyint",
+]
+class ColumnDefinition(BaseModel):
+    name: str
+    type: DbldatagenBasicType | None = None
+    primary: bool = False
+    options: dict[str, Any] | None = {}
+    nullable: bool | None = False
+    omit: bool | None = False
+    baseColumn: str | None = "id"
+    baseColumnType: str | None = "auto"
+
+    @root_validator()
+    def check_model_constraints(cls, values: dict[str, Any]) -> dict[str, Any]:
+        """
+        Validates constraints across the entire model after individual fields are processed.
+        """
+        is_primary = values.get("primary")
+        options = values.get("options", {})
+        name = values.get("name")
+        is_nullable = values.get("nullable")
+        column_type = values.get("type")
+
+        if is_primary:
+            if "min" in options or "max" in options:
+                raise ValueError(f"Primary column '{name}' cannot have min/max options.")
+
+            if is_nullable:
+                raise ValueError(f"Primary column '{name}' cannot be nullable.")
+
+            if column_type is None:
+                raise ValueError(f"Primary column '{name}' must have a type defined.")
+        return values

dbldatagen/spec/compat.py

@@ -3,11 +3,11 @@
 try:
     # This will succeed on environments with Pydantic V2.x
     # It imports the V1 API that is bundled within V2.
-    from pydantic.v1 import BaseModel, Field, validator, constr
+    from pydantic.v1 import BaseModel, Field, validator, constr, root_validator
 
 except ImportError:
     # This will be executed on environments with only Pydantic V1.x
-    from pydantic import BaseModel, Field, validator, constr, root_validator, field_validator
+    from pydantic import BaseModel, Field, validator, constr, root_validator
 
 # In your application code, do this:
 # from .compat import BaseModel

dbldatagen/spec/generator_spec.py

@@ -1,65 +1,21 @@
-from .compat import BaseModel, validator, root_validator, field_validator
-from typing import Dict, Optional, Union, Any, Literal, List
-import pandas as pd
-from IPython.display import display, HTML
-
-DbldatagenBasicType = Literal[
-    "string",
-    "int",
-    "long",
-    "float",
-    "double",
-    "decimal",
-    "boolean",
-    "date",
-    "timestamp",
-    "short",
-    "byte",
-    "binary",
-    "integer",
-    "bigint",
-    "tinyint",
-]
-
-class ColumnDefinition(BaseModel):
-    name: str
-    type: Optional[DbldatagenBasicType] = None
-    primary: bool = False
-    options: Optional[Dict[str, Any]] = {}
-    nullable: Optional[bool] = False
-    omit: Optional[bool] = False
-    baseColumn: Optional[str] = "id"
-    baseColumnType: Optional[str] = "auto"
-
-    @root_validator(skip_on_failure=True)
-    def check_model_constraints(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Validates constraints across the entire model after individual fields are processed.
-        """
-        is_primary = values.get("primary")
-        options = values.get("options", {})
-        name = values.get("name")
-        is_nullable = values.get("nullable")
-        column_type = values.get("type")
+from __future__ import annotations
 
-        if is_primary:
-            if "min" in options or "max" in options:
-                raise ValueError(f"Primary column '{name}' cannot have min/max options.")
+from typing import Any, Literal, Union
 
-            if is_nullable:
-                raise ValueError(f"Primary column '{name}' cannot be nullable.")
+import pandas as pd
+from IPython.display import HTML, display
+
+from dbldatagen.spec.column_spec import ColumnDefinition
 
-            if column_type is None:
-                raise ValueError(f"Primary column '{name}' must have a type defined.")
-        return values
+from .compat import BaseModel, validator
 
 
 class UCSchemaTarget(BaseModel):
     catalog: str
     schema_: str
     output_format: str = "delta"  # Default to delta for UC Schema
 
-    @field_validator("catalog", "schema_", mode="after")
+    @validator("catalog", "schema_")
     def validate_identifiers(cls, v):  # noqa: N805, pylint: disable=no-self-argument
         if not v.strip():
             raise ValueError("Identifier must be non-empty.")
@@ -76,7 +32,7 @@ class FilePathTarget(BaseModel):
     base_path: str
     output_format: Literal["csv", "parquet"]  # No default, must be specified
 
-    @field_validator("base_path", mode="after")
+    @validator("base_path")
     def validate_base_path(cls, v):  # noqa: N805, pylint: disable=no-self-argument
         if not v.strip():
             raise ValueError("base_path must be non-empty.")
@@ -88,16 +44,16 @@ def __str__(self):
 
 class TableDefinition(BaseModel):
     number_of_rows: int
-    partitions: Optional[int] = None
-    columns: List[ColumnDefinition]
+    partitions: int | None = None
+    columns: list[ColumnDefinition]
 
 
 class ValidationResult:
     """Container for validation results with errors and warnings."""
 
     def __init__(self) -> None:
-        self.errors: List[str] = []
-        self.warnings: List[str] = []
+        self.errors: list[str] = []
+        self.warnings: list[str] = []
 
     def add_error(self, message: str) -> None:
         """Add an error message."""
@@ -132,16 +88,16 @@ def __str__(self) -> str:
         return "\n".join(lines)
 
 class DatagenSpec(BaseModel):
-    tables: Dict[str, TableDefinition]
-    output_destination: Optional[Union[UCSchemaTarget, FilePathTarget]] = None # there is a abstraction, may be we can use that? talk to Greg
-    generator_options: Optional[Dict[str, Any]] = {}
-    intended_for_databricks: Optional[bool] = None # May be infered.
+    tables: dict[str, TableDefinition]
+    output_destination: Union[UCSchemaTarget, FilePathTarget] | None = None # there is a abstraction, may be we can use that? talk to Greg
+    generator_options: dict[str, Any] | None = {}
+    intended_for_databricks: bool | None = None # May be infered.
 
     def _check_circular_dependencies(
         self,
         table_name: str,
-        columns: List[ColumnDefinition]
-    ) -> List[str]:
+        columns: list[ColumnDefinition]
+    ) -> list[str]:
         """
         Check for circular dependencies in baseColumn references.
         Returns a list of error messages if circular dependencies are found.
@@ -284,17 +240,15 @@ def validate(self, strict: bool = True) -> ValidationResult:
                 "random", "randomSeed", "randomSeedMethod", "verbose",
                 "debug", "seedColumnName"
             ]
-            for key in self.generator_options.keys():
+            for key in self.generator_options:
                 if key not in known_options:
                     result.add_warning(
                         f"Unknown generator option: '{key}'. "
                         "This may be ignored during generation."
                     )
 
         # Now that all validations are complete, decide whether to raise
-        if strict and (result.errors or result.warnings):
-            raise ValueError(str(result))
-        elif not strict and result.errors:
+        if (strict and (result.errors or result.warnings)) or (not strict and result.errors):
             raise ValueError(str(result))
 
         return result

pyproject.toml

@@ -103,6 +103,7 @@ dependencies = [
     "jmespath>=0.10.0",
     "py4j>=0.10.9",
     "pickleshare>=0.7.5",
+    "ipython>=7.32.0",
 ]
 
 python="3.10"