Add return type, change names and refactor

etsap-TIMES · Nov 25, 2024 · c293838 · c293838
1 parent 3962730
commit c293838
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/xl2times/datatypes.py b/xl2times/datatypes.py
@@ -84,7 +84,7 @@ class DataModule(str, Enum):
     trade = "SuppXLS/Trades/ScenTrade_*.*"
 
     @classmethod
-    def is_module_type(cls, path):
+    def _find_type_(cls, path: str) -> Enum | None:
         for data_module in cls:
             if any(
                 PurePath(path.lower()).match(pattern.lower().strip())
@@ -94,8 +94,16 @@ def is_module_type(cls, path):
         return None
 
     @classmethod
-    def is_submodule_type(cls, path):
-        match cls.is_module_type(path):
+    def module_type(cls, path: str) -> str | None:
+        module_type = cls._find_type_(path)
+        if module_type:
+            return module_type.name
+        else:
+            return None
+
+    @classmethod
+    def submodule(cls, path: str) -> str | None:
+        match cls._find_type_(path):
             case DataModule.base | DataModule.subres:
                 if PurePath(path.lower()).match("*_trans.*"):
                     return "trans"
@@ -107,8 +115,8 @@ def is_submodule_type(cls, path):
                 return "main"
 
     @classmethod
-    def is_module_name(cls, path):
-        module_type = cls.is_module_type(path)
+    def module_name(cls, path: str) -> str | None:
+        module_type = cls._find_type_(path)
         match module_type:
             case DataModule.base | DataModule.sets | DataModule.lma | DataModule.demand | DataModule.trade | DataModule.syssettings:
                 return module_type.name

diff --git a/xl2times/transforms.py b/xl2times/transforms.py
@@ -359,9 +359,9 @@ def include_tables_source(
     def include_table_source(table: EmbeddedXlTable):
         df = table.dataframe.copy()
         df["source_filename"] = Path(table.filename).stem
-        df["data_module_type"] = DataModule.is_module_type(table.filename)
-        df["data_submodule"] = DataModule.is_submodule_type(table.filename)
-        df["data_module_name"] = DataModule.is_module_name(table.filename)
+        df["data_module_type"] = DataModule.module_type(table.filename)
+        df["data_submodule"] = DataModule.submodule(table.filename)
+        df["data_module_name"] = DataModule.module_name(table.filename)
         return replace(table, dataframe=df)
 
     return [include_table_source(table) for table in tables]