industrial-model 0.1.31__tar.gz → 0.1.33__tar.gz

{industrial_model-0.1.31 → industrial_model-0.1.33}/PKG-INFO

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

{industrial_model-0.1.31 → industrial_model-0.1.33}/industrial_model/cognite_adapters/__init__.py

@@ -92,7 +92,9 @@ class CogniteAdapter:
             next_cursor = query_result.cursors.get(view_external_id)
             data.extend(page_result)
 
-            last_page = len(page_result) < statement.limit_ or not next_cursor
+            last_page = (
+                len(page_result) < statement.get_values().limit or not next_cursor
+            )
             next_cursor_ = None if last_page else next_cursor
             cognite_query.cursors = {view_external_id: next_cursor_}
 

{industrial_model-0.1.31 → industrial_model-0.1.33}/industrial_model/cognite_adapters/aggregation_mapper.py

@@ -43,29 +43,33 @@ class AggregationMapper:
 
         root_view = self._view_mapper.get_view(root_node)
 
+        statement_values = statement.get_values()
         filters_ = (
-            filters.And(*self._filter_mapper.map(statement.where_clauses, root_view))
-            if statement.where_clauses
+            filters.And(
+                *self._filter_mapper.map(statement_values.where_clauses, root_view)
+            )
+            if statement_values.where_clauses
             else None
         )
+        aggregation_property = statement_values.aggregation_property.property
         metric_aggregation: MetricAggregation | None = None
         if statement.aggregate == "avg":
-            metric_aggregation = Avg(statement.aggregation_property.property)
+            metric_aggregation = Avg(aggregation_property)
         elif statement.aggregate == "min":
-            metric_aggregation = Min(statement.aggregation_property.property)
+            metric_aggregation = Min(aggregation_property)
         elif statement.aggregate == "max":
-            metric_aggregation = Max(statement.aggregation_property.property)
+            metric_aggregation = Max(aggregation_property)
         elif statement.aggregate == "sum":
-            metric_aggregation = Sum(statement.aggregation_property.property)
+            metric_aggregation = Sum(aggregation_property)
         elif statement.aggregate == "count":
-            metric_aggregation = Count(statement.aggregation_property.property)
+            metric_aggregation = Count(aggregation_property)
 
         if metric_aggregation is None:
-            raise ValueError(f"Unsupported aggregate function: {statement.aggregate_}")
+            raise ValueError(f"Unsupported aggregate function: {statement.aggregate}")
 
         group_by_columns = (
-            [prop.property for prop in statement.group_by_properties]
-            if statement.group_by_properties is not None
+            [prop.property for prop in statement_values.group_by_properties]
+            if statement_values.group_by_properties is not None
             else statement.entity.get_group_by_fields()
         )
         return AggregationQuery(
@@ -73,5 +77,5 @@ class AggregationMapper:
             metric_aggregation=metric_aggregation,
             filters=filters_,
             group_by_columns=group_by_columns,
-            limit=statement.limit_,
+            limit=statement_values.limit,
         )

{industrial_model-0.1.31 → industrial_model-0.1.33}/industrial_model/cognite_adapters/filter_mapper.py

@@ -2,12 +2,17 @@ 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
 from industrial_model.statements import (
     BoolExpression,
+    Column,
     Expression,
     LeafExpression,
 )
@@ -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.31 → industrial_model-0.1.33}/industrial_model/cognite_adapters/optimizer.py

@@ -30,7 +30,7 @@ class QueryOptimizer:
         if not instance_spaces and not instance_spaces_prefix:
             return
 
-        if self._has_space_filter(statement.where_clauses):
+        if self._has_space_filter(statement.get_values().where_clauses):
             return
 
         filter_spaces = (

{industrial_model-0.1.31 → industrial_model-0.1.33}/industrial_model/cognite_adapters/query_mapper.py

@@ -45,27 +45,36 @@ class QueryMapper:
 
         filters_: list[filters.Filter] = [filters.HasData(views=[root_view_id])]
 
-        filters_.extend(self._filter_mapper.map(statement.where_clauses, root_view))
+        statement_values = statement.get_values()
+        filters_.extend(
+            self._filter_mapper.map(statement_values.where_clauses, root_view)
+        )
 
         with_: dict[str, ResultSetExpression] = {
             root_node: NodeResultSetExpression(
                 filter=filters.And(*filters_),
-                sort=self._sort_mapper.map(statement.sort_clauses, root_view),
-                limit=statement.limit_,
+                sort=self._sort_mapper.map(statement_values.sort_clauses, root_view),
+                limit=statement_values.limit,
             )
         }
         select_: dict[str, Select] = {}
 
         relations = get_schema_properties(statement.entity, NESTED_SEP, root_node)
 
+        edge_filters = self._filter_mapper.map_edges(
+            statement_values.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)
 
         return CogniteQuery(
-            with_=with_, select=select_, cursors={root_node: statement.cursor_}
+            with_=with_,
+            select=select_,
+            cursors={root_node: statement_values.cursor},
         )
 
     def _get_select(self, view_id: ViewId, properties: list[str]) -> Select:
@@ -80,6 +89,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 +111,7 @@ class QueryMapper:
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     relations_to_include,
+                    edge_filters,
                     with_,
                     select_,
                 )
@@ -121,6 +132,7 @@ class QueryMapper:
                     property_key,
                     self._view_mapper.get_view(property.source.external_id),
                     relations_to_include,
+                    edge_filters,
                     with_,
                     select_,
                 )
@@ -139,10 +151,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 +175,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.31 → industrial_model-0.1.33}/industrial_model/cognite_adapters/search_mapper.py

@@ -31,20 +31,20 @@ class SearchMapper:
 
     def map(self, statement: SearchStatement[TViewInstance]) -> SearchQuery:
         root_node = statement.entity.get_view_external_id()
-
+        statement_values = statement.get_values()
         root_view = self._view_mapper.get_view(root_node)
 
-        filters_ = self._filter_mapper.map(statement.where_clauses, root_view)
+        filters_ = self._filter_mapper.map(statement_values.where_clauses, root_view)
 
-        sort_clauses = self._sort_mapper.map(statement.sort_clauses, root_view)
+        sort_clauses = self._sort_mapper.map(statement_values.sort_clauses, root_view)
         for item in sort_clauses:
             item.nulls_first = None
 
         return SearchQuery(
             view=root_view,
             filter=filters.And(*filters_) if filters_ else None,
-            query=statement.query,
-            query_properties=statement.query_properties,
-            limit=statement.limit_,
+            query=statement_values.query,
+            query_properties=statement_values.query_properties,
+            limit=statement_values.limit,
             sort=sort_clauses,
         )

{industrial_model-0.1.31 → industrial_model-0.1.33}/industrial_model/engines/engine.py

@@ -59,7 +59,7 @@ class Engine:
         statement: Statement[TViewInstance],
         validation_mode: ValidationMode = "raiseOnError",
     ) -> list[TViewInstance]:
-        if statement.cursor_:
+        if statement.get_values().cursor:
             raise ValueError("Cursor should be none when querying all pages")
 
         data, _ = self._cognite_adapter.query(statement, True)

{industrial_model-0.1.31 → industrial_model-0.1.33}/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.31 → industrial_model-0.1.33}/industrial_model/statements/__init__.py

@@ -23,22 +23,36 @@ def _create_column(property: str | Column | Any) -> Column:
 
 
 @dataclass
-class BaseStatement(Generic[T]):
-    entity: type[T] = field(init=True)
+class BaseStatementValues:
     where_clauses: list[Expression] = field(init=False, default_factory=list)
+    where_edge_clauses: list[tuple[Column, list[Expression]]] = field(
+        init=False, default_factory=list
+    )
     sort_clauses: list[tuple[Column, SORT_DIRECTION]] = field(
         init=False, default_factory=list
     )
-    limit_: int = field(init=False, default=DEFAULT_LIMIT)
+    limit: int = field(init=False, default=DEFAULT_LIMIT)
+    cursor: str | None = field(init=False, default=None)
+
+    query: str | None = field(init=False, default=None)
+    query_properties: list[str] | None = field(init=False, default=None)
+
+
+@dataclass
+class BaseStatement(Generic[T]):
+    entity: type[T] = field(init=True)
+    _values: BaseStatementValues = field(
+        init=False, default_factory=lambda: BaseStatementValues()
+    )
 
     def where(self, *expressions: bool | Expression) -> Self:
         for expression in expressions:
             assert isinstance(expression, Expression)
-            self.where_clauses.append(expression)
+            self._values.where_clauses.append(expression)
         return self
 
     def limit(self, limit: int) -> Self:
-        self.limit_ = limit
+        self._values.limit = limit
         return self
 
     def asc(self, property: str | Column | Any) -> Self:
@@ -48,7 +62,7 @@ class BaseStatement(Generic[T]):
         return self.sort(property, "descending")
 
     def sort(self, property: str | Column | Any, direction: SORT_DIRECTION) -> Self:
-        self.sort_clauses.append(
+        self._values.sort_clauses.append(
             (
                 _create_column(property),
                 direction,
@@ -56,28 +70,45 @@ class BaseStatement(Generic[T]):
         )
         return self
 
+    def get_values(self) -> BaseStatementValues:
+        return self._values
+
 
 @dataclass
 class Statement(BaseStatement[T]):
-    cursor_: str | None = field(init=False, default=None)
-
     def cursor(self, cursor: str | None) -> Self:
-        self.cursor_ = cursor
+        self._values.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._values.where_edge_clauses.append(
+            (
+                _create_column(property),
+                expressions_,
+            )
+        )
         return self
 
 
 @dataclass
 class SearchStatement(BaseStatement[T]):
-    query: str | None = field(init=False, default=None)
-    query_properties: list[str] | None = field(init=False, default=None)
-
     def query_by(
         self,
         query: str,
         query_properties: list[Column] | list[str] | list[Any] | None = None,
     ) -> Self:
-        self.query = query
-        self.query_properties = (
+        self._values.query = query
+        self._values.query_properties = (
             [
                 prop
                 if isinstance(prop, str)
@@ -92,34 +123,44 @@ class SearchStatement(BaseStatement[T]):
         return self
 
 
+@dataclass
+class AggregationStatementValues:
+    group_by_properties: list[Column] | None = field(init=False, default=None)
+    aggregation_property: Column = field(init=False, default=Column("externalId"))
+    where_clauses: list[Expression] = field(init=False, default_factory=list)
+    limit: int = field(init=False, default=-1)
+
+
 @dataclass
 class AggregationStatement(Generic[T]):
     entity: type[T] = field(init=True)
     aggregate: AggregateTypes = field(init=True)
 
-    group_by_properties: list[Column] | None = field(init=False, default=None)
-    aggregation_property: Column = field(init=False, default=Column("externalId"))
-    where_clauses: list[Expression] = field(init=False, default_factory=list)
-    limit_: int = field(init=False, default=-1)
+    _values: AggregationStatementValues = field(
+        init=False, default_factory=lambda: AggregationStatementValues()
+    )
 
     def group_by(self, *property: str | Column | Any) -> Self:
-        self.group_by_properties = [_create_column(prop) for prop in property]
+        self._values.group_by_properties = [_create_column(prop) for prop in property]
         return self
 
     def aggregate_by(self, property: str | Column | Any) -> Self:
-        self.aggregation_property = _create_column(property)
+        self._values.aggregation_property = _create_column(property)
         return self
 
     def where(self, *expressions: bool | Expression) -> Self:
         for expression in expressions:
             assert isinstance(expression, Expression)
-            self.where_clauses.append(expression)
+            self._values.where_clauses.append(expression)
         return self
 
     def limit(self, limit: int) -> Self:
-        self.limit_ = limit
+        self._values.limit = limit
         return self
 
+    def get_values(self) -> AggregationStatementValues:
+        return self._values
+
 
 def select(entity: type[T]) -> Statement[T]:
     return Statement(entity)

{industrial_model-0.1.31 → industrial_model-0.1.33}/pyproject.toml

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