Fix generation & mypy.

Deltares · Mar 14, 2024 · 2a8de52 · 2a8de52
1 parent bfa1a19
commit 2a8de52
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 32 deletions.
diff --git a/core/src/schema.jl b/core/src/schema.jl
@@ -7,12 +7,16 @@
 @schema "ribasim.basin.profile" BasinProfile
 @schema "ribasim.basin.state" BasinState
 @schema "ribasim.basin.subgrid" BasinSubgrid
+@schema "ribasim.basin.boundaryconcentration" BasinBoundaryconcentration
+@schema "ribasim.basin.concentration" BasinConcentration
 @schema "ribasim.terminal.static" TerminalStatic
 @schema "ribasim.fractionalflow.static" FractionalFlowStatic
 @schema "ribasim.flowboundary.static" FlowBoundaryStatic
 @schema "ribasim.flowboundary.time" FlowBoundaryTime
+@schema "ribasim.flowboundary.concentration" FlowBoundaryConcentration
 @schema "ribasim.levelboundary.static" LevelBoundaryStatic
 @schema "ribasim.levelboundary.time" LevelBoundaryTime
+@schema "ribasim.levelboundary.concentration" LevelBoundaryConcentration
 @schema "ribasim.linearresistance.static" LinearResistanceStatic
 @schema "ribasim.manningresistance.static" ManningResistanceStatic
 @schema "ribasim.pidcontrol.static" PidControlStatic
@@ -96,6 +100,22 @@ end
     urban_runoff::Union{Missing, Float64}
 end
 
+@version BasinBoundaryconcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    drainage_concentration::Union{Missing, Float64}
+    precipitation_concentration::Union{Missing, Float64}
+    urban_runoff_concentration::Union{Missing, Float64}
+end
+
+@version BasinConcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    concentration::Float64
+end
+
 @version BasinProfileV1 begin
     node_id::Int32
     area::Float64
@@ -132,6 +152,13 @@ end
     level::Float64
 end
 
+@version LevelBoundaryConcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    concentration::Float64
+end
+
 @version FlowBoundaryStaticV1 begin
     node_id::Int32
     active::Union{Missing, Bool}
@@ -144,6 +171,13 @@ end
     flow_rate::Float64
 end
 
+@version FlowBoundaryConcentrationV1 begin
+    node_id::Int32
+    time::DateTime
+    substance::String
+    concentration::Float64
+end
+
 @version LinearResistanceStaticV1 begin
     node_id::Int32
     active::Union{Missing, Bool}

diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
@@ -11,6 +11,22 @@ class Config:
         coerce = True
 
 
+class BasinBoundaryconcentrationSchema(_BaseSchema):
+    node_id: Series[int] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    drainage_concentration: Series[float] = pa.Field(nullable=True)
+    precipitation_concentration: Series[float] = pa.Field(nullable=True)
+    urban_runoff_concentration: Series[float] = pa.Field(nullable=True)
+
+
+class BasinConcentrationSchema(_BaseSchema):
+    node_id: Series[int] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=False)
+
+
 class BasinProfileSchema(_BaseSchema):
     node_id: Series[int] = pa.Field(nullable=False, default=0)
     area: Series[float] = pa.Field(nullable=False)
@@ -48,23 +64,6 @@ class BasinTimeSchema(_BaseSchema):
     urban_runoff: Series[float] = pa.Field(nullable=True)
 
 
-# No capital for concentration to get Basin / boundaryconcentration
-class BasinBoundaryconcentrationSchema(_BaseSchema):
-    node_id: Series[int] = pa.Field(nullable=False, default=0)
-    time: Series[Timestamp] = pa.Field(nullable=False)
-    substance: Series[str] = pa.Field(nullable=False)
-    drainage_concentration: Series[float] = pa.Field(nullable=True)
-    precipitation_concentration: Series[float] = pa.Field(nullable=True)
-    urban_runoff_concentration: Series[float] = pa.Field(nullable=True)
-
-
-class BasinConcentrationSchema(_BaseSchema):
-    node_id: Series[int] = pa.Field(nullable=False, default=0)
-    time: Series[Timestamp] = pa.Field(nullable=False)
-    substance: Series[str] = pa.Field(nullable=False)
-    concentration: Series[float] = pa.Field(nullable=False)
-
-
 class DiscreteControlConditionSchema(_BaseSchema):
     node_id: Series[int] = pa.Field(nullable=False, default=0)
     listen_node_type: Series[str] = pa.Field(nullable=False)
@@ -80,6 +79,13 @@ class DiscreteControlLogicSchema(_BaseSchema):
     control_state: Series[str] = pa.Field(nullable=False)
 
 
+class FlowBoundaryConcentrationSchema(_BaseSchema):
+    node_id: Series[int] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=False)
+
+
 class FlowBoundaryStaticSchema(_BaseSchema):
     node_id: Series[int] = pa.Field(nullable=False, default=0)
     active: Series[pa.BOOL] = pa.Field(nullable=True)
@@ -92,19 +98,19 @@ class FlowBoundaryTimeSchema(_BaseSchema):
     flow_rate: Series[float] = pa.Field(nullable=False)
 
 
-class FlowBoundaryConcentrationSchema(_BaseSchema):
-    node_id: Series[int] = pa.Field(nullable=False, default=0)
-    time: Series[Timestamp] = pa.Field(nullable=False)
-    substance: Series[str] = pa.Field(nullable=False)
-    concentration: Series[float] = pa.Field(nullable=False)
-
-
 class FractionalFlowStaticSchema(_BaseSchema):
     node_id: Series[int] = pa.Field(nullable=False, default=0)
     fraction: Series[float] = pa.Field(nullable=False)
     control_state: Series[str] = pa.Field(nullable=True)
 
 
+class LevelBoundaryConcentrationSchema(_BaseSchema):
+    node_id: Series[int] = pa.Field(nullable=False, default=0)
+    time: Series[Timestamp] = pa.Field(nullable=False)
+    substance: Series[str] = pa.Field(nullable=False)
+    concentration: Series[float] = pa.Field(nullable=False)
+
+
 class LevelBoundaryStaticSchema(_BaseSchema):
     node_id: Series[int] = pa.Field(nullable=False, default=0)
     active: Series[pa.BOOL] = pa.Field(nullable=True)
@@ -117,13 +123,6 @@ class LevelBoundaryTimeSchema(_BaseSchema):
     level: Series[float] = pa.Field(nullable=False)
 
 
-class LevelBoundaryConcentrationSchema(_BaseSchema):
-    node_id: Series[int] = pa.Field(nullable=False, default=0)
-    time: Series[Timestamp] = pa.Field(nullable=False)
-    substance: Series[str] = pa.Field(nullable=False)
-    concentration: Series[float] = pa.Field(nullable=False)
-
-
 class LevelDemandStaticSchema(_BaseSchema):
     node_id: Series[int] = pa.Field(nullable=False, default=0)
     min_level: Series[float] = pa.Field(nullable=False)

diff --git a/python/ribasim_testmodels/ribasim_testmodels/basic.py b/python/ribasim_testmodels/ribasim_testmodels/basic.py
@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from pathlib import Path
 from typing import Any
 
@@ -116,7 +117,7 @@ def basic_model() -> ribasim.Model:
     model.pump.add(Node(7, Point(4.0, 1.0)), [pump.Static(flow_rate=[0.5 / 3600])])
 
     # Setup flow boundary
-    flow_boundary_data = [
+    flow_boundary_data: Sequence[TableModel[Any]] = [
         flow_boundary.Static(flow_rate=[1e-4]),
         flow_boundary.Concentration(
             time="2020-01-01 00:00:00", substance=["Tracer"], concentration=[1.0]