Skip to content

Commit

Permalink
Make more python functions private (#1702)
Browse files Browse the repository at this point in the history
Fixes #1651
  • Loading branch information
SouthEndMusic authored Aug 8, 2024
1 parent 672c5b2 commit 142150b
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 19 deletions.
14 changes: 11 additions & 3 deletions python/ribasim/ribasim/geometry/edge.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ def add(
)
self.df.index.name = "fid"

def get_where_edge_type(self, edge_type: str) -> NDArray[np.bool_]:
def _get_where_edge_type(self, edge_type: str) -> NDArray[np.bool_]:
assert self.df is not None
return (self.df.edge_type == edge_type).to_numpy()

Expand All @@ -118,6 +118,14 @@ def sort(self):
self.df.sort_index(inplace=True)

def plot(self, **kwargs) -> Axes:
"""Plot the edges of the model.
Parameters
----------
**kwargs : Dict
Supported: 'ax', 'color_flow', 'color_control'
"""

assert self.df is not None
kwargs = kwargs.copy() # Avoid side-effects
ax = kwargs.get("ax", None)
Expand All @@ -141,8 +149,8 @@ def plot(self, **kwargs) -> Axes:
kwargs_control["color"] = color_control
kwargs_control["label"] = "Control edge"

where_flow = self.get_where_edge_type("flow")
where_control = self.get_where_edge_type("control")
where_flow = self._get_where_edge_type("flow")
where_control = self._get_where_edge_type("control")

if not self.df[where_flow].empty:
self.df[where_flow].plot(**kwargs_flow)
Expand Down
22 changes: 11 additions & 11 deletions python/ribasim/ribasim/input_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ class BaseModel(PydanticBaseModel):
)

@classmethod
def fields(cls) -> list[str]:
def _fields(cls) -> list[str]:
"""Return the names of the fields contained in the Model."""
return list(cls.model_fields.keys())

Expand All @@ -117,7 +117,7 @@ class FileModel(BaseModel, ABC):

@model_validator(mode="before")
@classmethod
def check_filepath(cls, value: Any) -> Any:
def _check_filepath(cls, value: Any) -> Any:
# Enable initialization with a Path.
if isinstance(value, dict):
# Pydantic Model init requires a dict
Expand Down Expand Up @@ -173,7 +173,7 @@ class TableModel(FileModel, Generic[TableT]):

@field_validator("df")
@classmethod
def check_extra_columns(cls, v: DataFrame[TableT]):
def _check_extra_columns(cls, v: DataFrame[TableT]):
"""Allow only extra columns with `meta_` prefix."""
if isinstance(v, (pd.DataFrame, gpd.GeoDataFrame)):
for colname in v.columns:
Expand All @@ -187,7 +187,7 @@ def check_extra_columns(cls, v: DataFrame[TableT]):
return v

@model_serializer
def set_model(self) -> str | None:
def _set_model(self) -> str | None:
return str(self.filepath.name) if self.filepath is not None else None

@classmethod
Expand All @@ -213,14 +213,14 @@ def tablename(cls) -> str:

@model_validator(mode="before")
@classmethod
def check_dataframe(cls, value: Any) -> Any:
def _check_dataframe(cls, value: Any) -> Any:
# Enable initialization with a DataFrame.
if isinstance(value, pd.DataFrame | gpd.GeoDataFrame):
value = {"df": value}

return value

def node_ids(self) -> set[int]:
def _node_ids(self) -> set[int]:
node_ids: set[int] = set()
if self.df is not None and "node_id" in self.df.columns:
node_ids.update(self.df["node_id"])
Expand Down Expand Up @@ -394,7 +394,7 @@ class ChildModel(BaseModel):
_parent_field: str | None = None

@model_validator(mode="after")
def check_parent(self) -> "ChildModel":
def _check_parent(self) -> "ChildModel":
if self._parent is not None:
self._parent.model_fields_set.update({self._parent_field})
return self
Expand Down Expand Up @@ -432,7 +432,7 @@ def _layername(cls, field: str) -> str:
return f"{cls.get_input_type()}{delimiter}{field}"

def _tables(self) -> Generator[TableModel[Any], Any, None]:
for key in self.fields():
for key in self._fields():
attr = getattr(self, key)
if (
isinstance(attr, TableModel)
Expand All @@ -441,10 +441,10 @@ def _tables(self) -> Generator[TableModel[Any], Any, None]:
):
yield attr

def node_ids(self) -> set[int]:
def _node_ids(self) -> set[int]:
node_ids: set[int] = set()
for table in self._tables():
node_ids.update(table.node_ids())
node_ids.update(table._node_ids())
return node_ids

def _save(self, directory: DirectoryPath, input_dir: DirectoryPath):
Expand All @@ -457,7 +457,7 @@ def _repr_content(self) -> str:
Skip "empty" attributes: when the dataframe of a TableModel is None.
"""
content = []
for field in self.fields():
for field in self._fields():
attr = getattr(self, field)
if isinstance(attr, TableModel):
if attr.df is not None:
Expand Down
10 changes: 5 additions & 5 deletions python/ribasim/ribasim/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ class Model(FileModel):
edge: EdgeTable = Field(default_factory=EdgeTable)

@model_validator(mode="after")
def set_node_parent(self) -> "Model":
def _set_node_parent(self) -> "Model":
for (
k,
v,
Expand All @@ -108,14 +108,14 @@ def set_node_parent(self) -> "Model":
return self

@model_validator(mode="after")
def ensure_edge_table_is_present(self) -> "Model":
def _ensure_edge_table_is_present(self) -> "Model":
if self.edge.df is None:
self.edge.df = GeoDataFrame[EdgeSchema]()
self.edge.df.set_geometry("geometry", inplace=True, crs=self.crs)
return self

@field_serializer("input_dir", "results_dir")
def serialize_path(self, path: Path) -> str:
def _serialize_path(self, path: Path) -> str:
return str(path)

def model_post_init(self, __context: Any) -> None:
Expand All @@ -132,7 +132,7 @@ def __repr__(self) -> str:
"""
content = ["ribasim.Model("]
INDENT = " "
for field in self.fields():
for field in self._fields():
attr = getattr(self, field)
if isinstance(attr, EdgeTable):
content.append(f"{INDENT}{field}=Edge(...),")
Expand Down Expand Up @@ -289,7 +289,7 @@ def _load(cls, filepath: Path | None) -> dict[str, Any]:
return {}

@model_validator(mode="after")
def reset_contextvar(self) -> "Model":
def _reset_contextvar(self) -> "Model":
# Drop database info
context_file_loading.set({})
return self
Expand Down

0 comments on commit 142150b

Please sign in to comment.