industrial-model 0.1.30__tar.gz → 0.1.32__tar.gz

{industrial_model-0.1.30 → industrial_model-0.1.32}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: industrial-model
-Version: 0.1.30
+Version: 0.1.32
 Summary: Industrial Model ORM
 Author-email: Lucas Alves <lucasrosaalves@gmail.com>
 Classifier: Programming Language :: Python

{industrial_model-0.1.30 → industrial_model-0.1.32}/industrial_model/cognite_adapters/filter_mapper.py

@@ -2,7 +2,11 @@ from datetime import date, datetime
 from typing import Any
 
 import cognite.client.data_classes.filters as cdf_filters
-from cognite.client.data_classes.data_modeling import MappedProperty, View
+from cognite.client.data_classes.data_modeling import (
+    EdgeConnection,
+    MappedProperty,
+    View,
+)
 
 from industrial_model.cognite_adapters.utils import get_property_ref
 from industrial_model.models.entities import InstanceId
@@ -11,6 +15,7 @@ from industrial_model.statements import (
     Expression,
     LeafExpression,
 )
+from industrial_model.statements.expressions import Column
 from industrial_model.utils import datetime_to_ms_iso_timestamp
 
 from .view_mapper import ViewMapper
@@ -26,15 +31,37 @@ class FilterMapper:
         result: list[cdf_filters.Filter] = []
         for expression in expressions:
             if isinstance(expression, BoolExpression):
-                result.append(self.to_cdf_filter_bool(expression, root_view))
+                result.append(self._to_cdf_filter_bool(expression, root_view))
             elif isinstance(expression, LeafExpression):
-                result.append(self.to_cdf_filter_leaf(expression, root_view))
+                result.append(self._to_cdf_filter_leaf(expression, root_view))
             else:
                 cls_name = expression.__class__.__name__
                 raise ValueError(f"Expression not implemented {cls_name}")
         return result
 
-    def to_cdf_filter_bool(
+    def map_edges(
+        self,
+        edges_expressions: list[tuple[Column, list[Expression]]],
+        root_view: View,
+        nested_separator: str,
+    ) -> dict[str, list[cdf_filters.Filter]]:
+        result_dict: dict[str, list[cdf_filters.Filter]] = {}
+
+        for column, expressions in edges_expressions:
+            view_property = root_view.properties.get(column.property)
+            if not isinstance(view_property, EdgeConnection):
+                raise ValueError(f"Property {column.property} is not an edge")
+
+            filters = self.map(
+                expressions,
+                self._view_mapper.get_view(view_property.source.external_id),
+            )
+            result_key = root_view.external_id + nested_separator + column.property
+            result_dict.setdefault(result_key, []).extend(filters)
+
+        return result_dict
+
+    def _to_cdf_filter_bool(
         self, expression: BoolExpression, root_view: View
     ) -> cdf_filters.Filter:
         arguments = self.map(expression.filters, root_view)
@@ -48,7 +75,7 @@ class FilterMapper:
 
         raise NotImplementedError(f"Operator {self.operator} not implemented")
 
-    def to_cdf_filter_leaf(
+    def _to_cdf_filter_leaf(
         self,
         expression: LeafExpression,
         root_view: View,

{industrial_model-0.1.30 → industrial_model-0.1.32}/industrial_model/cognite_adapters/query_mapper.py

@@ -58,8 +58,12 @@ class QueryMapper:
 
         relations = get_schema_properties(statement.entity, NESTED_SEP, root_node)
 
+        edge_filters = self._filter_mapper.map_edges(
+            statement.where_edge_clauses, root_view, NESTED_SEP
+        )
+
         properties = self._include_statements(
-            root_node, root_view, relations, with_, select_
+            root_node, root_view, relations, edge_filters, with_, select_
         )
 
         select_[root_node] = self._get_select(root_view_id, properties)
@@ -80,6 +84,7 @@ class QueryMapper:
         key: str,
         view: View,
         relations_to_include: list[str] | None,
+        edge_filters: dict[str, list[filters.Filter]],
         with_: dict[str, ResultSetExpression],
         select_: dict[str, Select],
     ) -> list[str]:
@@ -101,6 +106,7 @@ class QueryMapper:
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     relations_to_include,
+                    edge_filters,
                     with_,
                     select_,
                 )
@@ -121,6 +127,7 @@ class QueryMapper:
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     relations_to_include,
+                    edge_filters,
                     with_,
                     select_,
                 )
@@ -139,10 +146,16 @@ class QueryMapper:
             elif isinstance(property, EdgeConnection) and property.source:
                 edge_property_key = f"{property_key}{NESTED_SEP}{EDGE_MARKER}"
 
+                edge_filter = edge_filters.get(property_key)
+
                 with_[edge_property_key] = EdgeResultSetExpression(
                     from_=key,
                     max_distance=1,
-                    filter=filters.Equals(["edge", "type"], property.type.dump()),
+                    filter=filters.Equals(
+                        ["edge", "type"],
+                        property.type.dump(),
+                    ),
+                    node_filter=filters.And(*edge_filter) if edge_filter else None,
                     direction=property.direction,
                     limit=MAX_LIMIT,
                 )
@@ -157,6 +170,7 @@ class QueryMapper:
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     relations_to_include,
+                    edge_filters,
                     with_,
                     select_,
                 )

{industrial_model-0.1.30 → industrial_model-0.1.32}/industrial_model/cognite_adapters/query_result_mapper.py

@@ -10,6 +10,10 @@ from cognite.client.data_classes.data_modeling import (
     NodeList,
     View,
 )
+from cognite.client.data_classes.data_modeling.data_types import (
+    ListablePropertyType,
+)
+from cognite.client.data_classes.data_modeling.instances import PropertyValue
 from cognite.client.data_classes.data_modeling.views import (
     MultiReverseDirectRelation,
     SingleReverseDirectRelation,
@@ -112,13 +116,13 @@ class QueryResultMapper:
                 ):
                     continue
 
-                element_key: tuple[str, str] = (
-                    (element.get("space", ""), element.get("externalId", ""))
-                    if isinstance(element, dict)
-                    else (node.space, node.external_id)
-                )
+                element_keys = self._get_element_keys(node, element)
 
-                node_entries = mapping_nodes.get(element_key)
+                node_entries = [
+                    item
+                    for element_key in element_keys
+                    for item in mapping_nodes.get(element_key, [])
+                ]
                 if not node_entries:
                     if mapping_key in properties:
                         properties.pop(mapping_key)
@@ -126,7 +130,11 @@ class QueryResultMapper:
 
                 entry_data = self.nodes_to_dict(node_entries)
                 properties[mapping_key] = entry_data if is_list else entry_data[0]
-                edge_entries = mapping_edges.get(element_key)
+                edge_entries = [
+                    item
+                    for element_key in element_keys
+                    for item in mapping_edges.get(element_key, [])
+                ]
                 if edge_entries:
                     edges_mapping[mapping_key] = self._edges_to_model(edge_entries)
             properties["_edges"] = edges_mapping
@@ -137,6 +145,21 @@ class QueryResultMapper:
 
         return dict(result)
 
+    def _get_element_keys(
+        self, node: Node, element: PropertyValue | None
+    ) -> list[tuple[str, str]]:
+        if isinstance(element, dict):
+            return [(element.get("space", ""), element.get("externalId", ""))]
+
+        if isinstance(element, list):
+            return [
+                (item.get("space", ""), item.get("externalId", ""))
+                for item in element
+                if isinstance(item, dict)
+            ]
+
+        return [(node.space, node.external_id)]
+
     def _get_property_mappings(
         self,
         key: str,
@@ -159,7 +182,10 @@ class QueryResultMapper:
                     self._view_mapper.get_view(property.source.external_id),
                     query_result,
                 )
-                is_list = False
+                is_list = (
+                    isinstance(property.type, ListablePropertyType)
+                    and property.type.is_list
+                )
                 connection_type = ConnectionTypeEnum.DIRECT_RELATION
             elif isinstance(property, SingleReverseDirectRelation) and property.source:
                 nodes = self._map_node_property(

{industrial_model-0.1.30 → industrial_model-0.1.32}/industrial_model/models/entities.py

@@ -16,7 +16,7 @@ from industrial_model.statements import Column
 from .base import DBModelMetaclass, RootModel
 
 
-class InstanceId(RootModel):
+class InstanceId(RootModel, metaclass=DBModelMetaclass):
     external_id: str
     space: str
 
@@ -51,7 +51,7 @@ class ViewInstanceConfig(TypedDict, total=False):
     view_code: str | None
 
 
-class ViewInstance(InstanceId, metaclass=DBModelMetaclass):
+class ViewInstance(InstanceId):
     view_config: ClassVar[ViewInstanceConfig] = ViewInstanceConfig()
 
     _edges: dict[str, list[EdgeContainer]] = PrivateAttr(default_factory=dict)

{industrial_model-0.1.30 → industrial_model-0.1.32}/industrial_model/statements/__init__.py

@@ -60,11 +60,33 @@ class BaseStatement(Generic[T]):
 @dataclass
 class Statement(BaseStatement[T]):
     cursor_: str | None = field(init=False, default=None)
+    where_edge_clauses: list[tuple[Column, list[Expression]]] = field(
+        init=False, default_factory=list
+    )
 
     def cursor(self, cursor: str | None) -> Self:
         self.cursor_ = cursor
         return self
 
+    def where_edge(
+        self, property: str | Column | Any, *expressions: bool | Expression
+    ) -> Self:
+        if not expressions:
+            return self
+
+        expressions_: list[Expression] = []
+        for expression in expressions:
+            assert isinstance(expression, Expression)
+            expressions_.append(expression)
+
+        self.where_edge_clauses.append(
+            (
+                _create_column(property),
+                expressions_,
+            )
+        )
+        return self
+
 
 @dataclass
 class SearchStatement(BaseStatement[T]):

{industrial_model-0.1.30 → industrial_model-0.1.32}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "industrial-model"
-version = "0.1.30"
+version = "0.1.32"
 description = "Industrial Model ORM"
 readme = "README.md"
 requires-python = ">=3.11"