Skip to content

Commit

Permalink
enum
Browse files Browse the repository at this point in the history
  • Loading branch information
daanvaningen committed Oct 16, 2023
1 parent 5aff299 commit d7c3b58
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 17 deletions.
20 changes: 12 additions & 8 deletions threedigrid/admin/gridresultadmin.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@
PumpsResultsMixin,
)
from threedigrid.admin.structure_controls.models import (
STRUCTURE_CONTROL_TYPES,
StructureControl,
StructureControlTypes,
)
from threedigrid.orm.models import Model

Expand Down Expand Up @@ -320,15 +320,15 @@ def __init__(

@property
def table_control(self) -> "_GridH5NestedStructureControl":
return _GridH5NestedStructureControl(self, "table_control")
return _GridH5NestedStructureControl(self, StructureControlTypes.table_control)

@property
def memory_control(self) -> "_GridH5NestedStructureControl":
return _GridH5NestedStructureControl(self, "memory_control")
return _GridH5NestedStructureControl(self, StructureControlTypes.memory_control)

@property
def timed_control(self) -> "_GridH5NestedStructureControl":
return _GridH5NestedStructureControl(self, "timed_control")
return _GridH5NestedStructureControl(self, StructureControlTypes.timed_control)

def get_source_table(self, action_type, grid_id):
"""Get source_table and source_table_id based on action_type and grid_id"""
Expand All @@ -343,7 +343,11 @@ def get_source_table(self, action_type, grid_id):


class _GridH5NestedStructureControl:
def __init__(self, structure_control: GridH5StructureControl, control_type: str):
def __init__(
self,
structure_control: GridH5StructureControl,
control_type: StructureControlTypes,
):
"""
:param structure_control: GridH5StructureControl(GridH5ResultAdmin)
:param control_type: str [table_control, memory_control, timed_control]
Expand All @@ -352,11 +356,11 @@ def __init__(self, structure_control: GridH5StructureControl, control_type: str)
(usually structure_control_actions_3di.nc)
:param file_modus: modus in which to open the files
"""
if control_type not in STRUCTURE_CONTROL_TYPES:
raise ValueError(f"Unknown control type {control_type}")
if control_type not in StructureControlTypes.__members__.values():
raise ValueError(f"Unknown control type: {control_type}")

self.struct_control = structure_control
self.control_type = control_type
self.control_type: str = control_type.value

@property
def action_type(self) -> np.ndarray:
Expand Down
8 changes: 4 additions & 4 deletions threedigrid/admin/structure_controls/exporters.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
_GridH5NestedStructureControl,
GridH5StructureControl,
)
from threedigrid.admin.structure_controls.models import STRUCTURE_CONTROL_TYPES
from threedigrid.admin.structure_controls.models import StructureControlTypes


def structure_control_actions_to_csv(
Expand All @@ -26,9 +26,9 @@ def structure_control_actions_to_csv(
"is_active",
]
)
for control_type in STRUCTURE_CONTROL_TYPES:
for control_type in StructureControlTypes.__members__.values():
control_type_data: _GridH5NestedStructureControl = getattr(
structure_control, control_type
structure_control, control_type.name
)
for (
id,
Expand All @@ -52,7 +52,7 @@ def structure_control_actions_to_csv(
)
csv_writer.writerow(
[
control_type,
control_type.value,
id,
source_table,
source_table_id,
Expand Down
11 changes: 6 additions & 5 deletions threedigrid/admin/structure_controls/models.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
from enum import Enum
from typing import List

STRUCTURE_CONTROL_TYPES: List[str] = [
"table_control",
"memory_control",
"timed_control",
]

class StructureControlTypes(Enum):
table_control = "table_control"
memory_control = "memory_control"
timed_control = "timed_control"


class StructureControl:
Expand Down

0 comments on commit d7c3b58

Please sign in to comment.