Use union types. (#793)

Co-authored-by: Martijn Visser <[email protected]>

docs/contribute/addnode.qmd

@@ -143,7 +143,7 @@ class NewNodeType(TableModel):
     possible other schemas
     """
 
-    static: Optional[DataFrame[StaticSchema]] = None
+    static: DataFrame[StaticSchema] | None
     # possible other schemas
 
     model_config = ConfigDict(validate_assignment=True)

pixi.toml

@@ -77,7 +77,7 @@ tests = { depends_on = ["lint", "test-ribasim-python", "test-ribasim-core"] }
 generate-schema = { cmd = "julia --project=docs docs/gen_schema.jl", depends_on = [
     "instantiate-julia",
 ] }
-generate-python = "datamodel-codegen --use-title-as-name --use-double-quotes --disable-timestamp --use-default --strict-nullable --input-file-type=jsonschema --input docs/schema/root.schema.json --output python/ribasim/ribasim/models.py"
+generate-python = "datamodel-codegen --use-union-operator --use-title-as-name --use-double-quotes --disable-timestamp --use-default --strict-nullable --input-file-type=jsonschema --input docs/schema/root.schema.json --output python/ribasim/ribasim/models.py"
 codegen = { depends_on = ["generate-schema", "generate-python", "lint"] }
 # Publish
 build-ribasim-python-wheel = { cmd = "rm --recursive --force dist && python -m build && twine check dist/*", cwd = "python/ribasim" }

python/ribasim/ribasim/config.py

@@ -1,9 +1,6 @@
-# generated by datamodel-codegen:
-#   filename:  Config.schema.json
-
 from enum import Enum
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List
 
 from pydantic import Field
 
@@ -37,7 +34,7 @@
 
 
 class Allocation(BaseModel):
-    timestep: Optional[float] = None
+    timestep: float | None = None
     use_allocation: bool = False
 
 
@@ -50,7 +47,7 @@ class Results(BaseModel):
     basin: Path = Path("results/basin.arrow")
     flow: Path = Path("results/flow.arrow")
     control: Path = Path("results/control.arrow")
-    outstate: Optional[str] = None
+    outstate: str | None = None
     compression: Compression = Compression.zstd
     compression_level: int = 6
 
@@ -59,9 +56,9 @@ class Solver(BaseModel):
     algorithm: str = "QNDF"
     saveat: float | List[float] = []
     adaptive: bool = True
-    dt: Optional[float] = None
-    dtmin: Optional[float] = None
-    dtmax: Optional[float] = None
+    dt: float | None = None
+    dtmin: float | None = None
+    dtmax: float | None = None
     force_dtmin: bool = False
     abstol: float = 1e-06
     reltol: float = 1e-05

python/ribasim/ribasim/input_base.py

@@ -92,7 +92,7 @@ class FileModel(BaseModel, ABC):
 
     Attributes
     ----------
-        filepath (Optional[Path]):
+        filepath (Path | None):
             The path of this FileModel.
     """
 
@@ -255,7 +255,7 @@ def sort(self, sort_keys: List[str] = ["node_id"]):
     def tableschema(cls) -> TableT:
         """Retrieve Pandera Schema.
 
-        The type of the field `df` is known to always be an Optional[DataFrame[TableT]]]
+        The type of the field `df` is known to always be an DataFrame[TableT]]] | None
         """
         optionalfieldtype = cls.model_fields["df"].annotation
         fieldtype = optionalfieldtype.__args__[0]  # type: ignore

python/ribasim/ribasim/model.py

@@ -90,44 +90,58 @@ class Model(FileModel):
 
     Parameters
     ----------
-    node : Node
-        The ID, type and geometry of each node.
-    edge : Edge
-        How the nodes are connected.
+    starttime : datetime.datetime
+        Starting time of the simulation.
+    endtime : datetime.datetime
+        End time of the simulation.
+
+    update_timestep: float = 86400
+        The output time step of the simulation in seconds (default of 1 day)
+    relative_dir: str = "."
+        The relative directory of the input files.
+    input_dir: str = "."
+        The directory of the input files.
+    results_dir: str = "."
+        The directory of the results files.
+
+    network: Network
+        Class containing the topology (nodes and edges) of the model.
+
+    results: Results
+        Results configuration options.
+    solver: Solver
+        Solver configuration options.
+    logging: Logging
+        Logging configuration options.
+
+    allocation: Allocation
+        The allocation configuration.
     basin : Basin
         The waterbodies.
-    fractional_flow : Optional[FractionalFlow]
+    fractional_flow : FractionalFlow
         Split flows into fractions.
-    level_boundary : Optional[LevelBoundary]
+    level_boundary : LevelBoundary
         Boundary condition specifying the water level.
-    flow_boundary : Optional[FlowBoundary]
+    flow_boundary : FlowBoundary
         Boundary conditions specifying the flow.
-    linear_resistance: Optional[LinearResistance]
+    linear_resistance: LinearResistance
         Linear flow resistance.
-    manning_resistance : Optional[ManningResistance]
+    manning_resistance : ManningResistance
         Flow resistance based on the Manning formula.
-    tabulated_rating_curve : Optional[TabulatedRatingCurve]
+    tabulated_rating_curve : TabulatedRatingCurve
         Tabulated rating curve describing flow based on the upstream water level.
-    pump : Optional[Pump]
+    pump : Pump
         Prescribed flow rate from one basin to the other.
-    outlet : Optional[Outlet]
+    outlet : Outlet
         Prescribed flow rate from one basin to the other.
-    terminal : Optional[Terminal]
+    terminal : Terminal
         Water sink without state or properties.
-    discrete_control : Optional[DiscreteControl]
+    discrete_control : DiscreteControl
         Discrete control logic.
-    pid_control : Optional[PidControl]
+    pid_control : PidControl
         PID controller attempting to set the level of a basin to a desired value using a pump/outlet.
-    user : Optional[User]
+    user : User
         User node type with demand and priority.
-    starttime : Union[str, datetime.datetime]
-        Starting time of the simulation.
-    endtime : Union[str, datetime.datetime]
-        End time of the simulation.
-    solver : Optional[Solver]
-        Solver settings.
-    logging : Optional[logging]
-        Logging settings.
     """
 
     starttime: datetime.datetime

python/ribasim/ribasim/models.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Optional
 
 from pydantic import BaseModel, Field
 
@@ -48,7 +47,7 @@ class DiscreteControlCondition(BaseModel):
     listen_feature_id: int
     variable: str
     greater_than: float
-    look_ahead: Optional[float] = None
+    look_ahead: float | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
@@ -65,13 +64,13 @@ class Edge(BaseModel):
     from_node_id: int
     to_node_id: int
     edge_type: str
-    allocation_network_id: Optional[int] = None
+    allocation_network_id: int | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class FlowBoundaryStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     flow_rate: float
     remarks: str = Field("", description="a hack for pandera")
 
@@ -86,13 +85,13 @@ class FlowBoundaryTime(BaseModel):
 class FractionalFlowStatic(BaseModel):
     node_id: int
     fraction: float
-    control_state: Optional[str] = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class LevelBoundaryStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     level: float
     remarks: str = Field("", description="a hack for pandera")
 
@@ -106,51 +105,51 @@ class LevelBoundaryTime(BaseModel):
 
 class LinearResistanceStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     resistance: float
-    control_state: Optional[str] = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class ManningResistanceStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     length: float
     manning_n: float
     profile_width: float
     profile_slope: float
-    control_state: Optional[str] = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class Node(BaseModel):
     fid: int
     name: str
     type: str
-    allocation_network_id: Optional[int] = None
+    allocation_network_id: int | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class OutletStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     flow_rate: float
-    min_flow_rate: Optional[float] = None
-    max_flow_rate: Optional[float] = None
-    min_crest_level: Optional[float] = None
-    control_state: Optional[str] = None
+    min_flow_rate: float | None = None
+    max_flow_rate: float | None = None
+    min_crest_level: float | None = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class PidControlStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     listen_node_id: int
     target: float
     proportional: float
     integral: float
     derivative: float
-    control_state: Optional[str] = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
@@ -162,26 +161,26 @@ class PidControlTime(BaseModel):
     proportional: float
     integral: float
     derivative: float
-    control_state: Optional[str] = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class PumpStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     flow_rate: float
-    min_flow_rate: Optional[float] = None
-    max_flow_rate: Optional[float] = None
-    control_state: Optional[str] = None
+    min_flow_rate: float | None = None
+    max_flow_rate: float | None = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
 class TabulatedRatingCurveStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     level: float
     discharge: float
-    control_state: Optional[str] = None
+    control_state: str | None = None
     remarks: str = Field("", description="a hack for pandera")
 
 
@@ -200,7 +199,7 @@ class TerminalStatic(BaseModel):
 
 class UserStatic(BaseModel):
     node_id: int
-    active: Optional[bool] = None
+    active: bool | None = None
     demand: float
     return_factor: float
     min_level: float
@@ -219,27 +218,27 @@ class UserTime(BaseModel):
 
 
 class Root(BaseModel):
-    BasinProfile: Optional[BasinProfile] = None
-    BasinState: Optional[BasinState] = None
-    BasinStatic: Optional[BasinStatic] = None
-    BasinTime: Optional[BasinTime] = None
-    DiscreteControlCondition: Optional[DiscreteControlCondition] = None
-    DiscreteControlLogic: Optional[DiscreteControlLogic] = None
-    Edge: Optional[Edge] = None
-    FlowBoundaryStatic: Optional[FlowBoundaryStatic] = None
-    FlowBoundaryTime: Optional[FlowBoundaryTime] = None
-    FractionalFlowStatic: Optional[FractionalFlowStatic] = None
-    LevelBoundaryStatic: Optional[LevelBoundaryStatic] = None
-    LevelBoundaryTime: Optional[LevelBoundaryTime] = None
-    LinearResistanceStatic: Optional[LinearResistanceStatic] = None
-    ManningResistanceStatic: Optional[ManningResistanceStatic] = None
-    Node: Optional[Node] = None
-    OutletStatic: Optional[OutletStatic] = None
-    PidControlStatic: Optional[PidControlStatic] = None
-    PidControlTime: Optional[PidControlTime] = None
-    PumpStatic: Optional[PumpStatic] = None
-    TabulatedRatingCurveStatic: Optional[TabulatedRatingCurveStatic] = None
-    TabulatedRatingCurveTime: Optional[TabulatedRatingCurveTime] = None
-    TerminalStatic: Optional[TerminalStatic] = None
-    UserStatic: Optional[UserStatic] = None
-    UserTime: Optional[UserTime] = None
+    BasinProfile: BasinProfile | None
+    BasinState: BasinState | None
+    BasinStatic: BasinStatic | None
+    BasinTime: BasinTime | None
+    DiscreteControlCondition: DiscreteControlCondition | None
+    DiscreteControlLogic: DiscreteControlLogic | None
+    Edge: Edge | None
+    FlowBoundaryStatic: FlowBoundaryStatic | None
+    FlowBoundaryTime: FlowBoundaryTime | None
+    FractionalFlowStatic: FractionalFlowStatic | None
+    LevelBoundaryStatic: LevelBoundaryStatic | None
+    LevelBoundaryTime: LevelBoundaryTime | None
+    LinearResistanceStatic: LinearResistanceStatic | None
+    ManningResistanceStatic: ManningResistanceStatic | None
+    Node: Node | None
+    OutletStatic: OutletStatic | None
+    PidControlStatic: PidControlStatic | None
+    PidControlTime: PidControlTime | None
+    PumpStatic: PumpStatic | None
+    TabulatedRatingCurveStatic: TabulatedRatingCurveStatic | None
+    TabulatedRatingCurveTime: TabulatedRatingCurveTime | None
+    TerminalStatic: TerminalStatic | None
+    UserStatic: UserStatic | None
+    UserTime: UserTime | None

ruff.toml

@@ -1,5 +1,5 @@
 # See https://docs.astral.sh/ruff/rules/
-select = ["C4", "D2", "D3", "D4", "E", "F", "I", "NPY", "PD"]
+select = ["C4", "D2", "D3", "D4", "E", "F", "I", "NPY", "PD", "UP"]
 ignore = [
     "D202",
     "D205",
@@ -14,9 +14,7 @@ ignore = [
 ]
 fixable = ["I"]
 extend-include = ["*.ipynb"]
-exclude = [
-    "ribasim_qgis/tomllib/*"
-]
+exclude = ["ribasim_qgis/tomllib/*"]
 
 [pydocstyle]
 convention = "numpy"