Attribute view

deepfield/field/_misc/load_welltrack.py

@@ -5,10 +5,10 @@
 import pandas as pd
 
 
-def load_welltrack(data: DataType, binary_data: BinaryData, logger: logging.Logger) -> dict[str, pd.DataFrame] | None:
+def load_welltrack(data: DataType, binary_data: BinaryData, logger: logging.Logger) -> pd.DataFrame | None:
     _ = binary_data, logger
     section = 'SCHEDULE'
-    res: dict[str, pd.DataFrame] = {}
+    res: list[pd.DataFrame] = []
     if not section in data:
         return None
     for key, val in data[section]:
@@ -17,9 +17,7 @@ def load_welltrack(data: DataType, binary_data: BinaryData, logger: logging.Logg
             assert len(val) == 2
             assert isinstance(val[0], str)
             assert isinstance(val[1], pd.DataFrame)
-            res[cast(str, val[0])] = cast(pd.DataFrame, val[1])
-
+            res.append(cast(pd.DataFrame, val[1]).assign(WELL=cast(str, val[0])))
     if not res:
         return None
-    return res
-
+    return pd.concat(res)

deepfield/field/base_tree_node.py

@@ -1,45 +1,34 @@
 """TreeSegment components."""
 from __future__ import annotations
-from abc import ABC, abstractmethod
-from typing import Generic, TypeVar, override
+from typing import TypeVar
 from anytree import NodeMixin
 import pandas as pd
 
 from .base_component import BaseComponent
 
 T = TypeVar('T', dict[str, pd.DataFrame], pd.DataFrame)
 
-class NodeAttributeViewBase(ABC, Generic[T]):
+class NodeAttributeView():
     def __init__(self, att: str, key: str | None) -> None:
         self._att: str = att
         self._key: str | None = key
+
     def __get__(self, obj: BaseTreeNode, objtype=None) -> pd.DataFrame | None:
         _ = objtype
         name = obj.name
         assert isinstance(name, str)
         component = obj.root_component
         comp_att = getattr(component, self._att)
         if comp_att is None:
-            return comp_att
+            return None
         return self._get(comp_att, name)
 
-    @abstractmethod
-    def _get(self, att: T, name: str) -> pd.DataFrame:
-        pass
-
-    def __set__(self, obj: BaseTreeNode, value: pd.DataFrame):
-        raise NotImplementedError()
-
-class NodeAttributeViewDataFrame(NodeAttributeViewBase[pd.DataFrame]):
-    @override
     def _get(self, att: pd.DataFrame, name: str):
         assert self._key is not None
         return att[att[self._key] == name]
 
-class NodeAttributeViewDict(NodeAttributeViewBase[dict[str, pd.DataFrame]]):
-    @override
-    def _get(self, att: dict[str, pd.DataFrame], name: str):
-        return att[name]
+    def __set__(self, obj: BaseTreeNode, value: pd.DataFrame) -> None:
+        raise NotImplementedError()
 
 
 class BaseTreeNode(BaseComponent, NodeMixin):

deepfield/field/well_segment.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 
-from .base_tree_node import BaseTreeNode, NodeAttributeViewDataFrame, NodeAttributeViewDict
+from .base_tree_node import BaseTreeNode, NodeAttributeView
 
 
 class WellSegment(BaseTreeNode):
@@ -30,16 +30,16 @@ class WellSegment(BaseTreeNode):
         Node's full name from root.
     """
 
-    wconprod: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('WCONPROD', 'WELL')
-    wconinj: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('WCONINJ', 'WELL')
-    welspecs: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('WELSPECS', 'WELL')
-    welspecsl: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('WELSPECSL', 'WELL')
-    compdat: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('COMPDAT', 'WELL')
-    compdatl: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('COMPDATL', 'WELL')
-    compdatmd: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('COMPDATMD', 'WELL')
-    wefac: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('WEFAC', 'WELL')
-    results: NodeAttributeViewDataFrame = NodeAttributeViewDataFrame('RESULTS', 'WELL')
-    welltrack: NodeAttributeViewDict = NodeAttributeViewDict('WELLTRACK', None)
+    wconprod: NodeAttributeView = NodeAttributeView('WCONPROD', 'WELL')
+    wconinj: NodeAttributeView = NodeAttributeView('WCONINJ', 'WELL')
+    welspecs: NodeAttributeView = NodeAttributeView('WELSPECS', 'WELL')
+    welspecsl: NodeAttributeView = NodeAttributeView('WELSPECSL', 'WELL')
+    compdat: NodeAttributeView = NodeAttributeView('COMPDAT', 'WELL')
+    compdatl: NodeAttributeView = NodeAttributeView('COMPDATL', 'WELL')
+    compdatmd: NodeAttributeView = NodeAttributeView('COMPDATMD', 'WELL')
+    wefac: NodeAttributeView = NodeAttributeView('WEFAC', 'WELL')
+    results: NodeAttributeView = NodeAttributeView('RESULTS', 'WELL')
+    welltrack: NodeAttributeView = NodeAttributeView('WELLTRACK', 'WELL')
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)