industrial-model 0.1.31__py3-none-any.whl → 0.1.32__py3-none-any.whl

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/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/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/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.31.dist-info → industrial_model-0.1.32.dist-info}/METADATA

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

{industrial_model-0.1.31.dist-info → industrial_model-0.1.32.dist-info}/RECORD

@@ -5,10 +5,10 @@ industrial_model/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 industrial_model/utils.py,sha256=oh4AxwxXaWgIC2uolkCbvkgo0ququHB6yAPVIXy45Ts,663
 industrial_model/cognite_adapters/__init__.py,sha256=cKPW5y-wWiyWYLPDaNjhqc5zdsgdLc40vunIlz2Debg,6246
 industrial_model/cognite_adapters/aggregation_mapper.py,sha256=qglWUnyhgSPG005HzIRyjYkt9V2WKsab5y4ruwuKLcA,2503
-industrial_model/cognite_adapters/filter_mapper.py,sha256=R_OX8JOM33vzBVFvfo0iDmSJl5rhGjsTBLyKDoesHO0,4349
+industrial_model/cognite_adapters/filter_mapper.py,sha256=mo-SmUzeNYPKdO8Ed_Kz02RWXMRqsZP7VVjKPWQtn4g,5308
 industrial_model/cognite_adapters/models.py,sha256=2j2IS01uPkQEp9WdVk8seYzEqGcDdWFnpzXhusHB2zk,945
 industrial_model/cognite_adapters/optimizer.py,sha256=EQfCe-4mSeXhYa2kcqJerPXYsCRgMJPA1AIf1o1KNH0,2415
-industrial_model/cognite_adapters/query_mapper.py,sha256=QdFVkfmL4tWcCy_oypOluT0Fo68obPk2RjBPvVA9mUs,5950
+industrial_model/cognite_adapters/query_mapper.py,sha256=ui-FhF7ixsJU8jAFNotdZ8vfSKtnjzZYkvLNgMU27Jc,6467
 industrial_model/cognite_adapters/query_result_mapper.py,sha256=KaBqYr5l1t8sQMxEoX_F3J2prypKIPaDJxSzGF_fMB8,10203
 industrial_model/cognite_adapters/search_mapper.py,sha256=dh69Te8oiDIRPbEMTcof-ZJ4XD-xB5iiBHUR0UETzho,1504
 industrial_model/cognite_adapters/sort_mapper.py,sha256=RJUAYlZGXoYzK0PwX63cibRF_L-MUq9g2ZsC2EeNIF4,696
@@ -20,14 +20,14 @@ industrial_model/engines/async_engine.py,sha256=-sQv2vn93bBmTZOxY_C1r40YP7IMl1ND
 industrial_model/engines/engine.py,sha256=lECOpjN6fGvKt28b441isT_u4UNF3LeLiJH-zwHmkYw,3798
 industrial_model/models/__init__.py,sha256=AzJ0CyPK5PvUCX45FFtybl13tkukUvk2UAF_90s_LQ8,742
 industrial_model/models/base.py,sha256=iGhDjXqA5ULEQIFHtkMi7WYJl0nQq1wi8_zqOr-Ep78,1649
-industrial_model/models/entities.py,sha256=BC85P8THwvVmcNpYq2spJuFJtlK8Oegkmd0kx3CI_wU,4704
+industrial_model/models/entities.py,sha256=XugtGFCJrHnxR4UX7FWFumB2tNZGB7mNjM4skznnmdQ,4704
 industrial_model/models/schemas.py,sha256=EoLK9GYdS-0DQ98myTP3wOU9YpWIJOfDSLOYZUy3xEY,4559
 industrial_model/queries/__init__.py,sha256=lA83zOxMRgBgkseWJgK9kCr1vD8D8iSWs9NGGRnoYKk,355
 industrial_model/queries/models.py,sha256=VK69c4L0b0miPrKvOQBB8A01SPxZXYThrquv6gw5OGY,2544
 industrial_model/queries/params.py,sha256=50qY5BO5onLsXorhcv-7qCKhJaMO94UzhKLCmZKY55s,1667
 industrial_model/queries/utils.py,sha256=uP6PLh9IVHDK6J8x444zHWPmyV4PkxdLO-PMc6qWItc,1505
-industrial_model/statements/__init__.py,sha256=rjLRo2KoazHQaOpmPkxbI3_Nm8NCkJxjpuqow6IZVSc,4221
+industrial_model/statements/__init__.py,sha256=dNEFceuirobo3c9b8KB-3qgKihZ4yfnUj5U-1978eBM,4856
 industrial_model/statements/expressions.py,sha256=4ZZOcZroI5-4xRw4PXIRlufi0ARndE5zSbbxLDpR2Ec,4816
-industrial_model-0.1.31.dist-info/METADATA,sha256=BZq-MwEpDSndxkVUf5A6odc9ENMiUuOa_MnKjLTjcSo,6858
-industrial_model-0.1.31.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-industrial_model-0.1.31.dist-info/RECORD,,
+industrial_model-0.1.32.dist-info/METADATA,sha256=SgEMILIv7zqrEK1Wyb-l-74ixnL2s9KHvcUrOvvHeHo,6858
+industrial_model-0.1.32.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+industrial_model-0.1.32.dist-info/RECORD,,