From d7c3b58b3777c8322af42c3aab46ddbd5d9f33d3 Mon Sep 17 00:00:00 2001
From: daanvaningen <daanvaningen@gmail.com>
Date: Mon, 16 Oct 2023 10:13:19 +0200
Subject: [PATCH] enum

---
 threedigrid/admin/gridresultadmin.py          | 20 +++++++++++--------
 .../admin/structure_controls/exporters.py     |  8 ++++----
 .../admin/structure_controls/models.py        | 11 +++++-----
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/threedigrid/admin/gridresultadmin.py b/threedigrid/admin/gridresultadmin.py
index 9dbe0a4..41c6c2b 100644
--- a/threedigrid/admin/gridresultadmin.py
+++ b/threedigrid/admin/gridresultadmin.py
@@ -33,8 +33,8 @@
     PumpsResultsMixin,
 )
 from threedigrid.admin.structure_controls.models import (
-    STRUCTURE_CONTROL_TYPES,
     StructureControl,
+    StructureControlTypes,
 )
 from threedigrid.orm.models import Model
 
@@ -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"""
@@ -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]
@@ -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:
diff --git a/threedigrid/admin/structure_controls/exporters.py b/threedigrid/admin/structure_controls/exporters.py
index 8a9b975..5355436 100644
--- a/threedigrid/admin/structure_controls/exporters.py
+++ b/threedigrid/admin/structure_controls/exporters.py
@@ -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(
@@ -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,
@@ -52,7 +52,7 @@ def structure_control_actions_to_csv(
                 )
                 csv_writer.writerow(
                     [
-                        control_type,
+                        control_type.value,
                         id,
                         source_table,
                         source_table_id,
diff --git a/threedigrid/admin/structure_controls/models.py b/threedigrid/admin/structure_controls/models.py
index 62e8691..0a61ce6 100644
--- a/threedigrid/admin/structure_controls/models.py
+++ b/threedigrid/admin/structure_controls/models.py
@@ -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: