![](\"images/windturbine_model.png\")
\n"
]
},
{
@@ -63,7 +46,7 @@
"metadata": {},
"outputs": [],
"source": [
- "from equipment_unit import EquipmentUnitClient"
+ "from wind_turbine import WindTurbineClient"
]
},
{
@@ -73,7 +56,7 @@
"metadata": {},
"outputs": [],
"source": [
- "pygen = EquipmentUnitClient.from_toml(\"config.toml\")"
+ "pygen = WindTurbineClient.from_toml(\"config.toml\")"
]
},
{
@@ -89,12 +72,12 @@
"id": "a095350c-834d-4606-8582-b536ca8a9282",
"metadata": {},
"source": [
- "The `StartEndTime` edges are available on the `unit_procedure` class, and have a `list` method we can use to list and filter all edges of this type."
+ "The `Distance` edges are available on the `wind_turbine` class, and have a `list` method we can use to list and filter all edges of this type."
]
},
{
"cell_type": "code",
- "execution_count": 4,
+ "execution_count": 6,
"id": "c5f449e7-7daf-42fb-8417-73f70f212094",
"metadata": {},
"outputs": [
@@ -124,102 +107,94 @@
" | \n", - " | space | \n", - "external_id | \n", - "edge_type | \n", - "start_node | \n", - "end_node | \n", - "end_time | \n", - "start_time | \n", - "data_record | \n", - "
|---|---|---|---|---|---|---|---|---|
| 0 | \n", - "IntegrationTestsImmutable | \n", - "unit_procedure:Scott Patel:equipment_module:Mi... | \n", - "{'space': 'IntegrationTestsImmutable', 'extern... | \n", - "{'space': 'IntegrationTestsImmutable', 'extern... | \n", - "{'space': 'IntegrationTestsImmutable', 'extern... | \n", - "2023-11-07 18:58:39+00:00 | \n", - "2023-11-04 23:22:42+00:00 | \n", - "{'version': 1, 'last_updated_time': 2023-11-13... | \n", - "
| \n", - " | space | \n", - "external_id | \n", - "edge_type | \n", - "start_node | \n", - "end_node | \n", - "end_time | \n", - "start_time | \n", - "data_record | \n", - "
|---|---|---|---|---|---|---|---|---|
| 0 | \n", - "IntegrationTestsImmutable | \n", - "unit_procedure:Kelly Evans:equipment_module:Gr... | \n", - "{'space': 'IntegrationTestsImmutable', 'extern... | \n", - "{'space': 'IntegrationTestsImmutable', 'extern... | \n", - "{'space': 'IntegrationTestsImmutable', 'extern... | \n", - "2023-11-13 18:01:26+00:00 | \n", - "2023-11-12 13:56:38+00:00 | \n", - "{'version': 1, 'last_updated_time': 2023-11-13... | \n", - "
| \n", - " | value | \n", - "
|---|---|
| space | \n", - "IntegrationTestsImmutable | \n", - "
| external_id | \n", - "unit_procedure:Matthew Gonzalez | \n", - "
| data_record | \n", - "{'version': 1, 'last_updated_time': 2023-11-13... | \n", - "
| node_type | \n", - "None | \n", - "
| name | \n", - "Matthew Gonzalez | \n", - "
| type_ | \n", - "red | \n", - "
| work_orders | \n", - "None | \n", - "
| work_units | \n", - "[{'space': 'IntegrationTestsImmutable', 'exter... | \n", - "
| \n", - " | value | \n", - "||
|---|---|---|---|
| space | \n", - "IntegrationTestsImmutable | \n", - "||
| external_id | \n", - "equipment_module:Hannah Mcgee | \n", + "{'space': 'sp_wind', 'external_id': 'utsira_st... | \n", "|
| data_record | \n", - "{'version': 1, 'last_updated_time': 2023-11-13... | \n", - "||
| node_type | \n", - "None | \n", - "||
| description | \n", - "Box likely camera short part different half. L... | \n", - "||
| name | \n", - "Hannah Mcgee | \n", - "||
| sensor_value | \n", - "Hospital wear trouble part about question. Def... | \n", + "{'version': 1, 'last_updated_time': 2024-11-16... | \n", "|
| type_ | \n", - "orange | \n", + "distance | \n", + "1000.0 | \n", "
Call .list_full() to return a list of {nodes[0].title()} and -.list_{last_connection_name}() to return a list of {nodes[-1].title()}.
-""" - - -class NodeQueryCore(QueryCore[T_DomainModelList, T_DomainListEnd]): - _result_cls: ClassVar[type[DomainModel]] - - def list_full(self, limit: int = DEFAULT_QUERY_LIMIT) -> T_DomainModelList: - builder = self._create_query(limit, self._result_list_cls, return_step="first", try_reverse=True) - builder.execute_query(self._client, remove_not_connected=True) - return builder.unpack() - - def _list(self, limit: int = DEFAULT_QUERY_LIMIT) -> T_DomainListEnd: - builder = self._create_query(limit, cast(type[DomainModelList], self._result_list_cls_end), return_step="last") - for step in builder[:-1]: - step.select = None - builder.execute_query(self._client, remove_not_connected=False) - return builder.unpack() - - def _dump_yaml(self) -> str: - return self._create_query(DEFAULT_QUERY_LIMIT, self._result_list_cls)._dump_yaml() - - def _create_query( - self, - limit: int, - result_list_cls: type[DomainModelList], - return_step: Literal["first", "last"] | None = None, - try_reverse: bool = False, - ) -> DataClassQueryBuilder: - builder = DataClassQueryBuilder(result_list_cls, return_step=return_step) - from_: str | None = None - first: bool = True - is_last_reverse_list = False - for item in self._creation_path: - if is_last_reverse_list: - raise ValueError( - "Cannot traverse past reverse direct relation of list. " - "This is a limitation of the modeling implementation in your data model." - "To do this query, you need to reimplement the data model and use an edge to " - "implement this connection instead of a reverse direct relation" - ) - name = builder.create_name(from_) - max_retrieve_limit = limit if first else -1 - step: QueryStep - if isinstance(item, NodeQueryCore) and isinstance(item._expression, dm.query.NodeResultSetExpression): - step = NodeQueryStep( - name=name, - expression=item._expression, - result_cls=item._result_cls, - max_retrieve_limit=max_retrieve_limit, - connection_type=item._connection_type, - ) - step.expression.from_ = from_ - step.expression.filter = item._assemble_filter() - builder.append(step) - elif isinstance(item, NodeQueryCore) and isinstance(item._expression, dm.query.EdgeResultSetExpression): - edge_name = name - step = EdgeQueryStep(name=edge_name, expression=item._expression, max_retrieve_limit=max_retrieve_limit) - step.expression.from_ = from_ - builder.append(step) - - name = builder.create_name(edge_name) - node_step = NodeQueryStep( - name=name, - expression=dm.query.NodeResultSetExpression( - from_=edge_name, - filter=item._assemble_filter(), - ), - result_cls=item._result_cls, - ) - builder.append(node_step) - elif isinstance(item, EdgeQueryCore): - step = EdgeQueryStep( - name=name, - expression=cast(dm.query.EdgeResultSetExpression, item._expression), - result_cls=item._result_cls, - ) - step.expression.from_ = from_ - step.expression.filter = item._assemble_filter() - builder.append(step) - else: - raise TypeError(f"Unsupported query step type: {type(item._expression)}") - - is_last_reverse_list = item._connection_type == "reverse-list" - first = False - from_ = name - return builder - - -class EdgeQueryCore(QueryCore[T_DomainList, T_DomainListEnd]): - _result_cls: ClassVar[type[DomainRelation]] - - -@dataclass(frozen=True) -class ViewPropertyId(CogniteObject): - view: dm.ViewId - property: str - - @classmethod - def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None = None) -> "ViewPropertyId": - return cls( - view=dm.ViewId.load(resource["view"]), - property=resource["identifier"], - ) - - def dump(self, camel_case: bool = True) -> dict[str, Any]: - return { - "view": self.view.dump(camel_case=camel_case, include_type=False), - "identifier": self.property, - } - - -class QueryReducingBatchSize(UserWarning): - """Raised when a query is too large and the batch size must be reduced.""" - - ... - - -def chunker(sequence: Sequence, chunk_size: int) -> Iterator[Sequence]: - """ - Split a sequence into chunks of size chunk_size. - - Args: - sequence: The sequence to split. - chunk_size: The size of each chunk. - - Returns: - An iterator over the chunks. - - """ - for i in range(0, len(sequence), chunk_size): - yield sequence[i : i + chunk_size] - - -class QueryStep: - def __init__( - self, - name: str, - expression: dm.query.ResultSetExpression, - view_id: dm.ViewId | None = None, - max_retrieve_limit: int = -1, - select: dm.query.Select | None | type[_NotSetSentinel] = _NotSetSentinel, - raw_filter: dm.Filter | None = None, - connection_type: Literal["reverse-list"] | None = None, - view_property: ViewPropertyId | None = None, - selected_properties: list[str] | None = None, - ): - self.name = name - self.expression = expression - self.view_id = view_id - self.max_retrieve_limit = max_retrieve_limit - self.select: dm.query.Select | None - if select is _NotSetSentinel: - try: - self.select = self._default_select() - except NotImplementedError: - raise ValueError(f"You need to provide a select to instantiate a {type(self).__name__}") from None - else: - self.select = select # type: ignore[assignment] - self.raw_filter = raw_filter - self.connection_type = connection_type - self.view_property = view_property - self.selected_properties = selected_properties - self._max_retrieve_batch_limit = ACTUAL_INSTANCE_QUERY_LIMIT - self.cursor: str | None = None - self.total_retrieved: int = 0 - self.last_batch_count: int = 0 - self.results: list[Instance] = [] - - def _default_select(self) -> dm.query.Select: - if self.view_id is None: - return dm.query.Select() - else: - return dm.query.Select([dm.query.SourceSelector(self.view_id, ["*"])]) - - @property - def is_queryable(self) -> bool: - # We cannot query across reverse-list connections - return self.connection_type != "reverse-list" - - @property - def from_(self) -> str | None: - return self.expression.from_ - - @property - def is_single_direct_relation(self) -> bool: - return isinstance(self.expression, dm.query.NodeResultSetExpression) and self.expression.through is not None - - @property - def node_expression(self) -> dm.query.NodeResultSetExpression | None: - if isinstance(self.expression, dm.query.NodeResultSetExpression): - return self.expression - return None - - @property - def edge_expression(self) -> dm.query.EdgeResultSetExpression | None: - if isinstance(self.expression, dm.query.EdgeResultSetExpression): - return self.expression - return None - - @property - def node_results(self) -> Iterable[dm.Node]: - return (item for item in self.results if isinstance(item, dm.Node)) - - @property - def edge_results(self) -> Iterable[dm.Edge]: - return (item for item in self.results if isinstance(item, dm.Edge)) - - def update_expression_limit(self) -> None: - if self.is_unlimited: - self.expression.limit = self._max_retrieve_batch_limit - else: - self.expression.limit = max(min(INSTANCE_QUERY_LIMIT, self.max_retrieve_limit - self.total_retrieved), 0) - - def reduce_max_batch_limit(self) -> bool: - self._max_retrieve_batch_limit = max(1, self._max_retrieve_batch_limit // 2) - return self._max_retrieve_batch_limit > 1 - - @property - def is_unlimited(self) -> bool: - return self.max_retrieve_limit in {None, -1, math.inf} - - @property - def is_finished(self) -> bool: - return ( - (not self.is_unlimited and self.total_retrieved >= self.max_retrieve_limit) - or self.cursor is None - or self.last_batch_count == 0 - # Single direct relations are dependent on the parent node, - # so we assume that the parent node is the limiting factor. - or self.is_single_direct_relation - ) - - def count_total(self, cognite_client: CogniteClient) -> float: - if self.view_id is None: - raise ValueError("Cannot count total if select is not set") - - return cognite_client.data_modeling.instances.aggregate( - self.view_id, Count("externalId"), filter=self.raw_filter - ).value - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(name={self.name!r}, from={self.from_!r}, results={len(self.results)})" - - -class QueryBuilder(list, MutableSequence[QueryStep]): - """This is a helper class to build and execute a query. It is responsible for - doing the paging of the query and keeping track of the results.""" - - def __init__(self, steps: Collection[QueryStep] | None = None): - super().__init__(steps or []) - - def _reset(self): - for expression in self: - expression.total_retrieved = 0 - expression.cursor = None - expression.results = [] - - def _update_expression_limits(self) -> None: - for expression in self: - expression.update_expression_limit() - - def _build(self) -> tuple[dm.query.Query, list[QueryStep], set[str]]: - with_ = {step.name: step.expression for step in self if step.is_queryable} - select = {step.name: step.select for step in self if step.select is not None and step.is_queryable} - cursors = self._cursors - - step_by_name = {step.name: step for step in self} - search: list[QueryStep] = [] - temporary_select: set[str] = set() - for step in self: - if step.is_queryable: - continue - if step.node_expression is not None: - search.append(step) - # Ensure that select is set for the parent - if step.from_ in select or step.from_ is None: - continue - view_id = step_by_name[step.from_].view_id - if view_id is None: - continue - select[step.from_] = dm.query.Select([dm.query.SourceSelector(view_id, ["*"])]) - temporary_select.add(step.from_) - return dm.query.Query(with_=with_, select=select, cursors=cursors), search, temporary_select - - def _dump_yaml(self) -> str: - return self._build()[0].dump_yaml() - - @property - def _cursors(self) -> dict[str, str | None]: - return {expression.name: expression.cursor for expression in self if expression.is_queryable} - - def _update(self, batch: dm.query.QueryResult): - for expression in self: - if expression.name not in batch: - continue - expression.last_batch_count = len(batch[expression.name]) - expression.total_retrieved += expression.last_batch_count - expression.cursor = batch.cursors.get(expression.name) - expression.results.extend(batch[expression.name].data) - - @property - def _is_finished(self) -> bool: - return self[0].is_finished - - def _reduce_max_batch_limit(self) -> bool: - for expression in self: - if not expression.reduce_max_batch_limit(): - return False - return True - - def execute_query(self, client: CogniteClient, remove_not_connected: bool = False) -> dict[str, list[Instance]]: - self._reset() - query, to_search, temp_select = self._build() - - if not self: - raise ValueError("No query steps to execute") - - count: float | None = None - with suppress(ValueError, CogniteAPIError): - count = self[0].count_total(client) - - is_large_query = False - last_progress_print = 0 - nodes_per_second = 0.0 - while True: - self._update_expression_limits() - query.cursors = self._cursors - t0 = time.time() - try: - batch = client.data_modeling.instances.query(query) - except CogniteAPIError as e: - if e.code == 408: - # Too big query, try to reduce the limit - if self._reduce_max_batch_limit(): - continue - new_limit = self[0]._max_retrieve_batch_limit - warnings.warn( - f"Query is too large, reducing batch size to {new_limit:,}, and trying again", - QueryReducingBatchSize, - stacklevel=2, - ) - - raise e - - self._fetch_reverse_direct_relation_of_lists(client, to_search, batch) - - for name in temp_select: - batch.pop(name, None) - - last_execution_time = time.time() - t0 - - self._update(batch) - if self._is_finished: - break - - if count is None: - continue - # Estimate the number of nodes per second using exponential moving average - last_batch_nodes_per_second = len(batch[self[0].name]) / last_execution_time - if nodes_per_second == 0.0: - nodes_per_second = last_batch_nodes_per_second - else: - nodes_per_second = 0.1 * last_batch_nodes_per_second + 0.9 * nodes_per_second - # Estimate the time to completion - remaining_nodes = count - self[0].total_retrieved - remaining_time = remaining_nodes / nodes_per_second - - if is_large_query and (self[0].total_retrieved - last_progress_print) > PRINT_PROGRESS_PER_N_NODES: - estimate = datetime.timedelta(seconds=round(remaining_time, 0)) - print( - f"Progress: {self[0].total_retrieved:,}/{count:,} nodes retrieved. " - f"Estimated time to completion: {estimate}" - ) - last_progress_print = self[0].total_retrieved - - if is_large_query is False and remaining_time > MINIMUM_ESTIMATED_SECONDS_BEFORE_PRINT_PROGRESS: - is_large_query = True - print("Large query detected. Will print progress.") - - if remove_not_connected and len(self) > 1: - _QueryResultCleaner(self).clean() - - return {step.name: step.results for step in self} - - @staticmethod - def _fetch_reverse_direct_relation_of_lists( - client: CogniteClient, to_search: list[QueryStep], batch: dm.query.QueryResult - ) -> None: - """Reverse direct relations for lists are not supported by the query API. - This method fetches them separately.""" - for step in to_search: - if step.from_ is None or step.from_ not in batch: - continue - item_ids = [node.as_id() for node in batch[step.from_].data] - if not item_ids: - continue - - view_id = step.view_id - expression = step.node_expression - if view_id is None or expression is None: - raise ValueError( - "Invalid state of the query. Search should always be a node expression with view properties" - ) - if expression.through is None: - raise ValueError("Missing through set in a reverse-list query") - limit = SEARCH_LIMIT if step.is_unlimited else min(step.max_retrieve_limit, SEARCH_LIMIT) - - step_result = dm.NodeList[dm.Node]([]) - for item_ids_chunk in chunker(item_ids, IN_FILTER_CHUNK_SIZE): - is_items = dm.filters.In(view_id.as_property_ref(expression.through.property), item_ids_chunk) - is_selected = is_items if step.raw_filter is None else dm.filters.And(is_items, step.raw_filter) - - chunk_result = client.data_modeling.instances.search( - view_id, properties=None, filter=is_selected, limit=limit - ) - step_result.extend(chunk_result) - - batch[step.name] = dm.NodeListWithCursor(step_result, None) - return None - - def get_from(self) -> str | None: - if len(self) == 0: - return None - return self[-1].name - - def create_name(self, from_: str | None) -> str: - if from_ is None: - return "0" - return f"{from_}_{len(self)}" - - def append(self, __object: QueryStep, /) -> None: - # Extra validation to ensure all assumptions are met - if len(self) == 0: - if __object.from_ is not None: - raise ValueError("The first step should not have a 'from_' value") - else: - if __object.from_ is None: - raise ValueError("The 'from_' value should be set") - super().append(__object) - - def extend(self, __iterable: Iterable[QueryStep], /) -> None: - for item in __iterable: - self.append(item) - - # The implementations below are to get proper type hints - def __iter__(self) -> Iterator[QueryStep]: - return super().__iter__() - - @overload - def __getitem__(self, item: SupportsIndex) -> QueryStep: ... - - @overload - def __getitem__(self, item: slice) -> "QueryBuilder": ... - - def __getitem__(self, item: SupportsIndex | slice) -> "QueryStep | QueryBuilder": - value = super().__getitem__(item) - if isinstance(item, slice): - return QueryBuilder(value) # type: ignore[arg-type] - return cast(QueryStep, value) - - -class _QueryResultCleaner: - """Remove nodes and edges that are not connected through the entire query""" - - def __init__(self, steps: list[QueryStep]): - self._tree = self._create_tree(steps) - self._root = steps[0] - - @classmethod - def _create_tree(cls, steps: list[QueryStep]) -> dict[str, list[QueryStep]]: - tree: dict[str, list[QueryStep]] = defaultdict(list) - for step in steps: - if step.from_ is None: - continue - tree[step.from_].append(step) - return dict(tree) - - def clean(self) -> None: - self._clean(self._root) - - @staticmethod - def as_node_id(direct_relation: dm.DirectRelationReference | dict[str, str]) -> dm.NodeId: - if isinstance(direct_relation, dict): - return dm.NodeId(direct_relation["space"], direct_relation["externalId"]) - - return dm.NodeId(direct_relation.space, direct_relation.external_id) - - def _clean(self, step: QueryStep) -> tuple[set[dm.NodeId], str | None]: - if step.name not in self._tree: - # Leaf Node - direct_relation: str | None = None - if step.node_expression and (through := step.node_expression.through) is not None: - direct_relation = through.property - if step.node_expression.direction == "inwards": - return { - node_id for item in step.node_results for node_id in self._get_relations(item, direct_relation) - }, None - - return {item.as_id() for item in step.results}, direct_relation # type: ignore[attr-defined] - - expected_ids_by_property: dict[str | None, set[dm.NodeId]] = {} - for child in self._tree[step.name]: - child_ids, property_id = self._clean(child) - if property_id not in expected_ids_by_property: - expected_ids_by_property[property_id] = child_ids - else: - expected_ids_by_property[property_id] |= child_ids - - if step.node_expression is not None: - filtered_results: list[Instance] = [] - for node in step.node_results: - if self._is_connected_node(node, expected_ids_by_property): - filtered_results.append(node) - step.results = filtered_results - direct_relation = None if step.node_expression.through is None else step.node_expression.through.property - return {node.as_id() for node in step.node_results}, direct_relation - - if step.edge_expression: - if len(expected_ids_by_property) > 1 or None not in expected_ids_by_property: - raise RuntimeError(f"Invalid state of {type(self).__name__}") - expected_ids = expected_ids_by_property[None] - if step.edge_expression.direction == "outwards": - step.results = [edge for edge in step.edge_results if self.as_node_id(edge.end_node) in expected_ids] - return {self.as_node_id(edge.start_node) for edge in step.edge_results}, None - else: # inwards - step.results = [edge for edge in step.edge_results if self.as_node_id(edge.start_node) in expected_ids] - return {self.as_node_id(edge.end_node) for edge in step.edge_results}, None - - raise TypeError(f"Unsupported query step type: {type(step)}") - - @classmethod - def _is_connected_node(cls, node: dm.Node, expected_ids_by_property: dict[str | None, set[dm.NodeId]]) -> bool: - if not expected_ids_by_property: - return True - if None in expected_ids_by_property: - if node.as_id() in expected_ids_by_property[None]: - return True - if len(expected_ids_by_property) == 1: - return False - node_properties = next(iter(node.properties.values())) - for property_id, expected_ids in expected_ids_by_property.items(): - if property_id is None: - continue - value = node_properties.get(property_id) - if value is None: - continue - elif isinstance(value, list): - if {cls.as_node_id(item) for item in value if isinstance(item, dict)} & expected_ids: - return True - elif isinstance(value, dict) and cls.as_node_id(value) in expected_ids: - return True - return False - - @classmethod - def _get_relations(cls, node: dm.Node, property_id: str) -> Iterable[dm.NodeId]: - if property_id is None: - return {node.as_id()} - value = next(iter(node.properties.values())).get(property_id) - if isinstance(value, list): - return [cls.as_node_id(item) for item in value if isinstance(item, dict)] - elif isinstance(value, dict): - return [cls.as_node_id(value)] - return [] - - -class NodeQueryStep(QueryStep): - def __init__( - self, - name: str, - expression: dm.query.NodeResultSetExpression, - result_cls: type[DomainModel], - max_retrieve_limit: int = -1, - select: dm.query.Select | None | type[_NotSetSentinel] = _NotSetSentinel, - raw_filter: dm.Filter | None = None, - connection_type: Literal["reverse-list"] | None = None, - ): - self.result_cls = result_cls - super().__init__(name, expression, result_cls._view_id, max_retrieve_limit, select, raw_filter, connection_type) - - def unpack(self) -> dict[dm.NodeId | str, DomainModel]: - return { - ( - instance.as_id() if instance.space != DEFAULT_INSTANCE_SPACE else instance.external_id - ): self.result_cls.from_instance(instance) - for instance in cast(list[dm.Node], self.results) - } - - @property - def node_results(self) -> list[dm.Node]: - return cast(list[dm.Node], self.results) - - @property - def node_expression(self) -> dm.query.NodeResultSetExpression: - return cast(dm.query.NodeResultSetExpression, self.expression) - - -class EdgeQueryStep(QueryStep): - def __init__( - self, - name: str, - expression: dm.query.EdgeResultSetExpression, - result_cls: type[DomainRelation] | None = None, - max_retrieve_limit: int = -1, - select: dm.query.Select | None | type[_NotSetSentinel] = _NotSetSentinel, - raw_filter: dm.Filter | None = None, - ): - self.result_cls = result_cls - view_id = result_cls._view_id if result_cls is not None else None - super().__init__(name, expression, view_id, max_retrieve_limit, select, raw_filter, None) - - def unpack(self) -> dict[dm.NodeId, list[dm.Edge | DomainRelation]]: - output: dict[dm.NodeId, list[dm.Edge | DomainRelation]] = defaultdict(list) - for edge in cast(list[dm.Edge], self.results): - edge_source = edge.start_node if self.expression.direction == "outwards" else edge.end_node - value = self.result_cls.from_instance(edge) if self.result_cls is not None else edge - output[as_node_id(edge_source)].append(value) # type: ignore[arg-type] - return output - - @property - def edge_results(self) -> list[dm.Edge]: - return cast(list[dm.Edge], self.results) - - @property - def edge_expression(self) -> dm.query.EdgeResultSetExpression: - return cast(dm.query.EdgeResultSetExpression, self.expression) - - -class DataClassQueryBuilder(QueryBuilder, Generic[T_DomainModelList]): - """This is a helper class to build and execute a query. It is responsible for - doing the paging of the query and keeping track of the results.""" - - def __init__( - self, - result_cls: type[T_DomainModelList] | None, - steps: Collection[QueryStep] | None = None, - return_step: Literal["first", "last"] | None = None, - ): - super().__init__(steps or []) - self._result_list_cls = result_cls - self._return_step: Literal["first", "last"] | None = return_step - - def unpack(self) -> T_DomainModelList: - if self._result_list_cls is None: - raise ValueError("No result class set, unable to unpack results") - selected = [step for step in self if step.select is not None] - if len(selected) == 0: - return self._result_list_cls([]) - elif len(selected) == 1: - # Validated in the append method - if self._return_step == "first": - selected_step = cast(NodeQueryStep, self[0]) - elif self._return_step == "last": - selected_step = cast(NodeQueryStep, self[-1]) - else: - raise ValueError(f"Invalid return_step: {self._return_step}") - return self._result_list_cls(selected_step.unpack().values()) - # More than one step, we need to unpack the nodes and edges - nodes_by_from: dict[str | None, dict[dm.NodeId | str, DomainModel]] = defaultdict(dict) - edges_by_from: dict[str, dict[dm.NodeId, list[dm.Edge | DomainRelation]]] = defaultdict(dict) - for step in reversed(self): - # Validated in the append method - from_ = cast(str, step.from_) - if isinstance(step, EdgeQueryStep): - edges_by_from[from_].update(step.unpack()) - if step.name in nodes_by_from: - nodes_by_from[from_].update(nodes_by_from[step.name]) - del nodes_by_from[step.name] - elif isinstance(step, NodeQueryStep): - unpacked = step.unpack() - nodes_by_from[from_].update(unpacked) # type: ignore[arg-type] - if step.name in nodes_by_from or step.name in edges_by_from: - step.result_cls._update_connections( - unpacked, # type: ignore[arg-type] - nodes_by_from.get(step.name, {}), # type: ignore[arg-type] - edges_by_from.get(step.name, {}), - ) - if self._return_step == "first": - return self._result_list_cls(nodes_by_from[None].values()) - elif self._return_step == "last" and self[-1].from_ in nodes_by_from: - return self._result_list_cls(nodes_by_from[self[-1].from_].values()) - elif self._return_step == "last": - raise ValueError("Cannot return the last step when the last step is an edge query") - else: - raise ValueError(f"Invalid return_step: {self._return_step}") - - def append(self, __object: QueryStep, /) -> None: - # Extra validation to ensure all assumptions are met - if len(self) == 0: - if __object.from_ is not None: - raise ValueError("The first step should not have a 'from_' value") - if self._result_list_cls is None: - if self._return_step is None: - self._return_step = "first" - else: - if not isinstance(__object, NodeQueryStep): - raise ValueError("The first step should be a NodeQueryStep") - # If the first step is a NodeQueryStep, and matches the instance - # in the result_list_cls we can return the result from the first step - # Alternative is result_cls is not set, then we also assume that the first step - if self._return_step is None: - if __object.result_cls is self._result_list_cls._INSTANCE: - self._return_step = "first" - else: - # If not, we assume that the last step is the one we want to return - self._return_step = "last" - else: - if __object.from_ is None: - raise ValueError("The 'from_' value should be set") - super().append(__object) - - def extend(self, __iterable: Iterable[QueryStep], /) -> None: - for item in __iterable: - self.append(item) - - # The implementations below are to get proper type hints - def __iter__(self) -> Iterator[QueryStep]: - return super().__iter__() - - @overload - def __getitem__(self, item: SupportsIndex) -> QueryStep: ... - - @overload - def __getitem__(self, item: slice) -> DataClassQueryBuilder[T_DomainModelList]: ... - - def __getitem__(self, item: SupportsIndex | slice) -> QueryStep | DataClassQueryBuilder[T_DomainModelList]: - value = super().__getitem__(item) - if isinstance(item, slice): - return DataClassQueryBuilder(self._result_list_cls, value) # type: ignore[arg-type] - return cast(QueryStep, value) - - -T_QueryCore = TypeVar("T_QueryCore") - - -class Filtering(Generic[T_QueryCore], ABC): - def __init__(self, query: T_QueryCore, prop_path: list[str] | tuple[str, ...]): - self._query = query - self._prop_path = prop_path - self._filter: dm.Filter | None = None - - def _raise_if_filter_set(self): - if self._filter is not None: - raise ValueError("Filter has already been set") - - def _as_filter(self) -> dm.Filter | None: - return self._filter - - -class StringFilter(Filtering[T_QueryCore]): - def equals(self, value: str) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Equals(self._prop_path, value) - return self._query - - def prefix(self, prefix: str) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Prefix(self._prop_path, prefix) - return self._query - - def in_(self, values: list[str]) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.In(self._prop_path, values) - return self._query - - -class BooleanFilter(Filtering[T_QueryCore]): - def equals(self, value: bool) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Equals(self._prop_path, value) - return self._query - - -class IntFilter(Filtering[T_QueryCore]): - def range(self, gte: int | None, lte: int | None) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Range(self._prop_path, gte=gte, lte=lte) - return self._query - - -class FloatFilter(Filtering[T_QueryCore]): - def range(self, gte: float | None, lte: float | None) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Range(self._prop_path, gte=gte, lte=lte) - return self._query - - -class TimestampFilter(Filtering[T_QueryCore]): - def range(self, gte: datetime.datetime | None, lte: datetime.datetime | None) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Range( - self._prop_path, - gte=gte.isoformat(timespec="milliseconds") if gte else None, - lte=lte.isoformat(timespec="milliseconds") if lte else None, - ) - return self._query - - -class DateFilter(Filtering[T_QueryCore]): - def range(self, gte: datetime.date | None, lte: datetime.date | None) -> T_QueryCore: - self._raise_if_filter_set() - self._filter = dm.filters.Range( - self._prop_path, - gte=gte.isoformat() if gte else None, - lte=lte.isoformat() if lte else None, - ) - return self._query diff --git a/examples/equipment_unit/data_classes/_equipment_module.py b/examples/equipment_unit/data_classes/_equipment_module.py deleted file mode 100644 index 30b9a04bc..000000000 --- a/examples/equipment_unit/data_classes/_equipment_module.py +++ /dev/null @@ -1,396 +0,0 @@ -from __future__ import annotations - -import warnings -from collections.abc import Sequence -from typing import Any, ClassVar, Literal, no_type_check, Optional, Union - -from cognite.client import data_modeling as dm, CogniteClient -from cognite.client.data_classes import ( - TimeSeries as CogniteTimeSeries, - TimeSeriesWrite as CogniteTimeSeriesWrite, -) -from pydantic import Field -from pydantic import field_validator, model_validator - -from equipment_unit.data_classes._core import ( - DEFAULT_INSTANCE_SPACE, - DEFAULT_QUERY_LIMIT, - DataRecord, - DataRecordGraphQL, - DataRecordWrite, - DomainModel, - DomainModelWrite, - DomainModelWriteList, - DomainModelList, - DomainRelation, - DomainRelationWrite, - GraphQLCore, - ResourcesWrite, - FileMetadata, - FileMetadataWrite, - FileMetadataGraphQL, - TimeSeries, - TimeSeriesWrite, - TimeSeriesGraphQL, - T_DomainModelList, - as_direct_relation_reference, - as_instance_dict_id, - as_node_id, - as_pygen_node_id, - are_nodes_equal, - is_tuple_id, - select_best_node, - QueryCore, - NodeQueryCore, - StringFilter, -) - - -__all__ = [ - "EquipmentModule", - "EquipmentModuleWrite", - "EquipmentModuleApply", - "EquipmentModuleList", - "EquipmentModuleWriteList", - "EquipmentModuleApplyList", - "EquipmentModuleFields", - "EquipmentModuleTextFields", - "EquipmentModuleGraphQL", -] - - -EquipmentModuleTextFields = Literal["external_id", "description", "name", "sensor_value", "type_"] -EquipmentModuleFields = Literal["external_id", "description", "name", "sensor_value", "type_"] - -_EQUIPMENTMODULE_PROPERTIES_BY_FIELD = { - "external_id": "externalId", - "description": "description", - "name": "name", - "sensor_value": "sensor_value", - "type_": "type", -} - - -class EquipmentModuleGraphQL(GraphQLCore): - """This represents the reading version of equipment module, used - when data is retrieved from CDF using GraphQL. - - It is used when retrieving data from CDF using GraphQL. - - Args: - space: The space where the node is located. - external_id: The external id of the equipment module. - data_record: The data record of the equipment module node. - description: The description field. - name: The name field. - sensor_value: The sensor value field. - type_: The type field. - """ - - view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "EquipmentModule", "b1cd4bf14a7a33") - description: Optional[str] = None - name: Optional[str] = None - sensor_value: Optional[TimeSeriesGraphQL] = None - type_: Optional[str] = Field(None, alias="type") - - @model_validator(mode="before") - def parse_data_record(cls, values: Any) -> Any: - if not isinstance(values, dict): - return values - if "lastUpdatedTime" in values or "createdTime" in values: - values["dataRecord"] = DataRecordGraphQL( - created_time=values.pop("createdTime", None), - last_updated_time=values.pop("lastUpdatedTime", None), - ) - return values - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_read(self) -> EquipmentModule: - """Convert this GraphQL format of equipment module to the reading format.""" - if self.data_record is None: - raise ValueError("This object cannot be converted to a read format because it lacks a data record.") - return EquipmentModule( - space=self.space, - external_id=self.external_id, - data_record=DataRecord( - version=0, - last_updated_time=self.data_record.last_updated_time, - created_time=self.data_record.created_time, - ), - description=self.description, - name=self.name, - sensor_value=self.sensor_value.as_read() if self.sensor_value else None, - type_=self.type_, - ) - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_write(self) -> EquipmentModuleWrite: - """Convert this GraphQL format of equipment module to the writing format.""" - return EquipmentModuleWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=0), - description=self.description, - name=self.name, - sensor_value=self.sensor_value.as_write() if self.sensor_value else None, - type_=self.type_, - ) - - -class EquipmentModule(DomainModel): - """This represents the reading version of equipment module. - - It is used to when data is retrieved from CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the equipment module. - data_record: The data record of the equipment module node. - description: The description field. - name: The name field. - sensor_value: The sensor value field. - type_: The type field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "EquipmentModule", "b1cd4bf14a7a33") - - space: str = DEFAULT_INSTANCE_SPACE - node_type: Union[dm.DirectRelationReference, None] = None - description: Optional[str] = None - name: Optional[str] = None - sensor_value: Union[TimeSeries, str, None] = None - type_: Optional[str] = Field(None, alias="type") - - def as_write(self) -> EquipmentModuleWrite: - """Convert this read version of equipment module to the writing version.""" - return EquipmentModuleWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=self.data_record.version), - description=self.description, - name=self.name, - sensor_value=( - self.sensor_value.as_write() if isinstance(self.sensor_value, CogniteTimeSeries) else self.sensor_value - ), - type_=self.type_, - ) - - def as_apply(self) -> EquipmentModuleWrite: - """Convert this read version of equipment module to the writing version.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - -class EquipmentModuleWrite(DomainModelWrite): - """This represents the writing version of equipment module. - - It is used to when data is sent to CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the equipment module. - data_record: The data record of the equipment module node. - description: The description field. - name: The name field. - sensor_value: The sensor value field. - type_: The type field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "EquipmentModule", "b1cd4bf14a7a33") - - space: str = DEFAULT_INSTANCE_SPACE - node_type: Union[dm.DirectRelationReference, dm.NodeId, tuple[str, str], None] = None - description: Optional[str] = None - name: Optional[str] = None - sensor_value: Union[TimeSeriesWrite, str, None] = None - type_: Optional[str] = Field(None, alias="type") - - def _to_instances_write( - self, - cache: set[tuple[str, str]], - write_none: bool = False, - allow_version_increase: bool = False, - ) -> ResourcesWrite: - resources = ResourcesWrite() - if self.as_tuple_id() in cache: - return resources - - properties: dict[str, Any] = {} - - if self.description is not None or write_none: - properties["description"] = self.description - - if self.name is not None or write_none: - properties["name"] = self.name - - if self.sensor_value is not None or write_none: - properties["sensor_value"] = ( - self.sensor_value - if isinstance(self.sensor_value, str) or self.sensor_value is None - else self.sensor_value.external_id - ) - - if self.type_ is not None or write_none: - properties["type"] = self.type_ - - if properties: - this_node = dm.NodeApply( - space=self.space, - external_id=self.external_id, - existing_version=None if allow_version_increase else self.data_record.existing_version, - type=as_direct_relation_reference(self.node_type), - sources=[ - dm.NodeOrEdgeData( - source=self._view_id, - properties=properties, - ) - ], - ) - resources.nodes.append(this_node) - cache.add(self.as_tuple_id()) - - if isinstance(self.sensor_value, CogniteTimeSeriesWrite): - resources.time_series.append(self.sensor_value) - - return resources - - -class EquipmentModuleApply(EquipmentModuleWrite): - def __new__(cls, *args, **kwargs) -> EquipmentModuleApply: - warnings.warn( - "EquipmentModuleApply is deprecated and will be removed in v1.0. Use EquipmentModuleWrite instead." - "The motivation for this change is that Write is a more descriptive name for the writing version of the" - "EquipmentModule.", - UserWarning, - stacklevel=2, - ) - return super().__new__(cls) - - -class EquipmentModuleList(DomainModelList[EquipmentModule]): - """List of equipment modules in the read version.""" - - _INSTANCE = EquipmentModule - - def as_write(self) -> EquipmentModuleWriteList: - """Convert these read versions of equipment module to the writing versions.""" - return EquipmentModuleWriteList([node.as_write() for node in self.data]) - - def as_apply(self) -> EquipmentModuleWriteList: - """Convert these read versions of primitive nullable to the writing versions.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - -class EquipmentModuleWriteList(DomainModelWriteList[EquipmentModuleWrite]): - """List of equipment modules in the writing version.""" - - _INSTANCE = EquipmentModuleWrite - - -class EquipmentModuleApplyList(EquipmentModuleWriteList): ... - - -def _create_equipment_module_filter( - view_id: dm.ViewId, - description: str | list[str] | None = None, - description_prefix: str | None = None, - name: str | list[str] | None = None, - name_prefix: str | None = None, - type_: str | list[str] | None = None, - type_prefix: str | None = None, - external_id_prefix: str | None = None, - space: str | list[str] | None = None, - filter: dm.Filter | None = None, -) -> dm.Filter | None: - filters: list[dm.Filter] = [] - if isinstance(description, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("description"), value=description)) - if description and isinstance(description, list): - filters.append(dm.filters.In(view_id.as_property_ref("description"), values=description)) - if description_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("description"), value=description_prefix)) - if isinstance(name, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("name"), value=name)) - if name and isinstance(name, list): - filters.append(dm.filters.In(view_id.as_property_ref("name"), values=name)) - if name_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("name"), value=name_prefix)) - if isinstance(type_, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("type"), value=type_)) - if type_ and isinstance(type_, list): - filters.append(dm.filters.In(view_id.as_property_ref("type"), values=type_)) - if type_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("type"), value=type_prefix)) - if external_id_prefix is not None: - filters.append(dm.filters.Prefix(["node", "externalId"], value=external_id_prefix)) - if isinstance(space, str): - filters.append(dm.filters.Equals(["node", "space"], value=space)) - if space and isinstance(space, list): - filters.append(dm.filters.In(["node", "space"], values=space)) - if filter: - filters.append(filter) - return dm.filters.And(*filters) if filters else None - - -class _EquipmentModuleQuery(NodeQueryCore[T_DomainModelList, EquipmentModuleList]): - _view_id = EquipmentModule._view_id - _result_cls = EquipmentModule - _result_list_cls_end = EquipmentModuleList - - def __init__( - self, - created_types: set[type], - creation_path: list[QueryCore], - client: CogniteClient, - result_list_cls: type[T_DomainModelList], - expression: dm.query.ResultSetExpression | None = None, - connection_name: str | None = None, - connection_type: Literal["reverse-list"] | None = None, - reverse_expression: dm.query.ResultSetExpression | None = None, - ): - - super().__init__( - created_types, - creation_path, - client, - result_list_cls, - expression, - dm.filters.HasData(views=[self._view_id]), - connection_name, - connection_type, - reverse_expression, - ) - - self.space = StringFilter(self, ["node", "space"]) - self.external_id = StringFilter(self, ["node", "externalId"]) - self.description = StringFilter(self, self._view_id.as_property_ref("description")) - self.name = StringFilter(self, self._view_id.as_property_ref("name")) - self.type_ = StringFilter(self, self._view_id.as_property_ref("type")) - self._filter_classes.extend( - [ - self.space, - self.external_id, - self.description, - self.name, - self.type_, - ] - ) - - def list_equipment_module(self, limit: int = DEFAULT_QUERY_LIMIT) -> EquipmentModuleList: - return self._list(limit=limit) - - -class EquipmentModuleQuery(_EquipmentModuleQuery[EquipmentModuleList]): - def __init__(self, client: CogniteClient): - super().__init__(set(), [], client, EquipmentModuleList) diff --git a/examples/equipment_unit/data_classes/_start_end_time.py b/examples/equipment_unit/data_classes/_start_end_time.py deleted file mode 100644 index 80e9df41e..000000000 --- a/examples/equipment_unit/data_classes/_start_end_time.py +++ /dev/null @@ -1,428 +0,0 @@ -from __future__ import annotations - -import datetime -import warnings -from typing import Any, ClassVar, Literal, no_type_check, Optional, TYPE_CHECKING, Union - -from cognite.client import data_modeling as dm, CogniteClient - -from equipment_unit.data_classes._core import ( - DEFAULT_INSTANCE_SPACE, - DataRecord, - DataRecordWrite, - DomainModel, - DomainModelCore, - DomainModelWrite, - DomainRelation, - DomainRelationWrite, - DomainRelationList, - DomainRelationWriteList, - GraphQLCore, - ResourcesWrite, - DomainModelList, - T_DomainList, - EdgeQueryCore, - NodeQueryCore, - QueryCore, - StringFilter, - TimestampFilter, -) -from equipment_unit.data_classes._unit_procedure import UnitProcedureWrite -from equipment_unit.data_classes._equipment_module import EquipmentModule, EquipmentModuleGraphQL, EquipmentModuleWrite -from equipment_unit.data_classes._work_order import WorkOrder, WorkOrderGraphQL, WorkOrderWrite - -if TYPE_CHECKING: - from equipment_unit.data_classes._equipment_module import ( - EquipmentModule, - EquipmentModuleGraphQL, - EquipmentModuleWrite, - ) - from equipment_unit.data_classes._work_order import WorkOrder, WorkOrderGraphQL, WorkOrderWrite - - -__all__ = [ - "StartEndTime", - "StartEndTimeWrite", - "StartEndTimeApply", - "StartEndTimeList", - "StartEndTimeWriteList", - "StartEndTimeApplyList", - "StartEndTimeFields", -] - - -StartEndTimeTextFields = Literal["external_id",] -StartEndTimeFields = Literal["external_id", "end_time", "start_time"] -_STARTENDTIME_PROPERTIES_BY_FIELD = { - "external_id": "externalId", - "end_time": "end_time", - "start_time": "start_time", -} - - -class StartEndTimeGraphQL(GraphQLCore): - """This represents the reading version of start end time, used - when data is retrieved from CDF using GraphQL. - - It is used when retrieving data from CDF using GraphQL. - - Args: - space: The space where the node is located. - external_id: The external id of the start end time. - data_record: The data record of the start end time node. - end_node: The end node of this edge. - end_time: The end time field. - start_time: The start time field. - """ - - view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "StartEndTime", "d416e0ed98186b") - end_node: Union[EquipmentModuleGraphQL, WorkOrderGraphQL, None] = None - end_time: Optional[datetime.datetime] = None - start_time: Optional[datetime.datetime] = None - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_read(self) -> StartEndTime: - """Convert this GraphQL format of start end time to the reading format.""" - if self.data_record is None: - raise ValueError("This object cannot be converted to a read format because it lacks a data record.") - return StartEndTime( - space=self.space, - external_id=self.external_id, - data_record=DataRecord( - version=0, - last_updated_time=self.data_record.last_updated_time, - created_time=self.data_record.created_time, - ), - end_node=self.end_node.as_read() if isinstance(self.end_node, GraphQLCore) else self.end_node, - end_time=self.end_time, - start_time=self.start_time, - ) - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_write(self) -> StartEndTimeWrite: - """Convert this GraphQL format of start end time to the writing format.""" - return StartEndTimeWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=0), - end_node=self.end_node.as_write() if isinstance(self.end_node, DomainModel) else self.end_node, - end_time=self.end_time, - start_time=self.start_time, - ) - - -class StartEndTime(DomainRelation): - """This represents the reading version of start end time. - - It is used to when data is retrieved from CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the start end time. - data_record: The data record of the start end time edge. - end_node: The end node of this edge. - end_time: The end time field. - start_time: The start time field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "StartEndTime", "d416e0ed98186b") - space: str = DEFAULT_INSTANCE_SPACE - end_node: Union[EquipmentModule, WorkOrder, str, dm.NodeId] - end_time: Optional[datetime.datetime] = None - start_time: Optional[datetime.datetime] = None - - def as_write(self) -> StartEndTimeWrite: - """Convert this read version of start end time to the writing version.""" - return StartEndTimeWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=self.data_record.version), - end_node=self.end_node.as_write() if isinstance(self.end_node, DomainModel) else self.end_node, - end_time=self.end_time, - start_time=self.start_time, - ) - - def as_apply(self) -> StartEndTimeWrite: - """Convert this read version of start end time to the writing version.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - -class StartEndTimeWrite(DomainRelationWrite): - """This represents the writing version of start end time. - - It is used to when data is sent to CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the start end time. - data_record: The data record of the start end time edge. - end_node: The end node of this edge. - end_time: The end time field. - start_time: The start time field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "StartEndTime", "d416e0ed98186b") - space: str = DEFAULT_INSTANCE_SPACE - end_node: Union[EquipmentModuleWrite, WorkOrderWrite, str, dm.NodeId] - end_time: Optional[datetime.datetime] = None - start_time: Optional[datetime.datetime] = None - - def _to_instances_write( - self, - cache: set[tuple[str, str]], - start_node: DomainModelWrite, - edge_type: dm.DirectRelationReference, - write_none: bool = False, - allow_version_increase: bool = False, - ) -> ResourcesWrite: - resources = ResourcesWrite() - if self.external_id and (self.space, self.external_id) in cache: - return resources - - _validate_end_node(start_node, self.end_node) - - if isinstance(self.end_node, DomainModelWrite): - end_node = self.end_node.as_direct_reference() - elif isinstance(self.end_node, str): - end_node = dm.DirectRelationReference(self.space, self.end_node) - elif isinstance(self.end_node, dm.NodeId): - end_node = dm.DirectRelationReference(self.end_node.space, self.end_node.external_id) - else: - raise ValueError(f"Invalid type for equipment_module: {type(self.end_node)}") - - external_id = self.external_id or DomainRelationWrite.external_id_factory(start_node, self.end_node, edge_type) - - properties: dict[str, Any] = {} - - if self.end_time is not None or write_none: - properties["end_time"] = self.end_time.isoformat(timespec="milliseconds") if self.end_time else None - - if self.start_time is not None or write_none: - properties["start_time"] = self.start_time.isoformat(timespec="milliseconds") if self.start_time else None - - if properties: - this_edge = dm.EdgeApply( - space=self.space, - external_id=external_id, - type=edge_type, - start_node=start_node.as_direct_reference(), - end_node=end_node, - existing_version=None if allow_version_increase else self.data_record.existing_version, - sources=[ - dm.NodeOrEdgeData( - source=self._view_id, - properties=properties, - ) - ], - ) - resources.edges.append(this_edge) - cache.add((self.space, external_id)) - - if isinstance(self.end_node, DomainModelWrite): - other_resources = self.end_node._to_instances_write(cache) - resources.extend(other_resources) - - return resources - - -class StartEndTimeApply(StartEndTimeWrite): - def __new__(cls, *args, **kwargs) -> StartEndTimeApply: - warnings.warn( - "StartEndTimeApply is deprecated and will be removed in v1.0. Use StartEndTimeWrite instead." - "The motivation for this change is that Write is a more descriptive name for the writing version of the" - "StartEndTime.", - UserWarning, - stacklevel=2, - ) - return super().__new__(cls) - - -class StartEndTimeList(DomainRelationList[StartEndTime]): - """List of start end times in the reading version.""" - - _INSTANCE = StartEndTime - - def as_write(self) -> StartEndTimeWriteList: - """Convert this read version of start end time list to the writing version.""" - return StartEndTimeWriteList([edge.as_write() for edge in self]) - - def as_apply(self) -> StartEndTimeWriteList: - """Convert these read versions of start end time list to the writing versions.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - -class StartEndTimeWriteList(DomainRelationWriteList[StartEndTimeWrite]): - """List of start end times in the writing version.""" - - _INSTANCE = StartEndTimeWrite - - -class StartEndTimeApplyList(StartEndTimeWriteList): ... - - -def _create_start_end_time_filter( - edge_type: dm.DirectRelationReference, - view_id: dm.ViewId, - start_node: str | list[str] | dm.NodeId | list[dm.NodeId] | None = None, - start_node_space: str = DEFAULT_INSTANCE_SPACE, - end_node: str | list[str] | dm.NodeId | list[dm.NodeId] | None = None, - space_end_node: str = DEFAULT_INSTANCE_SPACE, - min_end_time: datetime.datetime | None = None, - max_end_time: datetime.datetime | None = None, - min_start_time: datetime.datetime | None = None, - max_start_time: datetime.datetime | None = None, - external_id_prefix: str | None = None, - space: str | list[str] | None = None, - filter: dm.Filter | None = None, -) -> dm.Filter: - filters: list[dm.Filter] = [ - dm.filters.Equals( - ["edge", "type"], - {"space": edge_type.space, "externalId": edge_type.external_id}, - ) - ] - if start_node and isinstance(start_node, str): - filters.append( - dm.filters.Equals(["edge", "startNode"], value={"space": start_node_space, "externalId": start_node}) - ) - if start_node and isinstance(start_node, dm.NodeId): - filters.append( - dm.filters.Equals( - ["edge", "startNode"], value=start_node.dump(camel_case=True, include_instance_type=False) - ) - ) - if start_node and isinstance(start_node, list): - filters.append( - dm.filters.In( - ["edge", "startNode"], - values=[ - ( - {"space": start_node_space, "externalId": ext_id} - if isinstance(ext_id, str) - else ext_id.dump(camel_case=True, include_instance_type=False) - ) - for ext_id in start_node - ], - ) - ) - if end_node and isinstance(end_node, str): - filters.append(dm.filters.Equals(["edge", "endNode"], value={"space": space_end_node, "externalId": end_node})) - if end_node and isinstance(end_node, dm.NodeId): - filters.append( - dm.filters.Equals(["edge", "endNode"], value=end_node.dump(camel_case=True, include_instance_type=False)) - ) - if end_node and isinstance(end_node, list): - filters.append( - dm.filters.In( - ["edge", "endNode"], - values=[ - ( - {"space": space_end_node, "externalId": ext_id} - if isinstance(ext_id, str) - else ext_id.dump(camel_case=True, include_instance_type=False) - ) - for ext_id in end_node - ], - ) - ) - if min_end_time is not None or max_end_time is not None: - filters.append( - dm.filters.Range( - view_id.as_property_ref("end_time"), - gte=min_end_time.isoformat(timespec="milliseconds") if min_end_time else None, - lte=max_end_time.isoformat(timespec="milliseconds") if max_end_time else None, - ) - ) - if min_start_time is not None or max_start_time is not None: - filters.append( - dm.filters.Range( - view_id.as_property_ref("start_time"), - gte=min_start_time.isoformat(timespec="milliseconds") if min_start_time else None, - lte=max_start_time.isoformat(timespec="milliseconds") if max_start_time else None, - ) - ) - if external_id_prefix is not None: - filters.append(dm.filters.Prefix(["edge", "externalId"], value=external_id_prefix)) - if isinstance(space, str): - filters.append(dm.filters.Equals(["edge", "space"], value=space)) - if space and isinstance(space, list): - filters.append(dm.filters.In(["edge", "space"], values=space)) - if filter: - filters.append(filter) - return dm.filters.And(*filters) - - -_EXPECTED_START_NODES_BY_END_NODE: dict[type[DomainModelWrite], set[type[DomainModelWrite]]] = { - EquipmentModuleWrite: {UnitProcedureWrite}, - WorkOrderWrite: {UnitProcedureWrite}, -} - - -def _validate_end_node( - start_node: DomainModelWrite, end_node: Union[EquipmentModuleWrite, WorkOrderWrite, str, dm.NodeId] -) -> None: - if isinstance(end_node, (str, dm.NodeId)): - # Nothing to validate - return - if type(end_node) not in _EXPECTED_START_NODES_BY_END_NODE: - raise ValueError( - f"Invalid end node type: {type(end_node)}. Should be one of {[t.__name__ for t in _EXPECTED_START_NODES_BY_END_NODE.keys()]}" - ) - if type(start_node) not in _EXPECTED_START_NODES_BY_END_NODE[type(end_node)]: - raise ValueError( - f"Invalid end node type: {type(end_node)}. Expected one of: {_EXPECTED_START_NODES_BY_END_NODE[type(end_node)]}" - ) - - -class _StartEndTimeQuery(EdgeQueryCore[T_DomainList, StartEndTimeList]): - _view_id = StartEndTime._view_id - _result_cls = StartEndTime - _result_list_cls_end = StartEndTimeList - - def __init__( - self, - created_types: set[type], - creation_path: list[QueryCore], - client: CogniteClient, - result_list_cls: type[T_DomainList], - end_node_cls: type[NodeQueryCore], - expression: dm.query.ResultSetExpression | None = None, - connection_name: str | None = None, - ): - from ._equipment_module import _EquipmentModuleQuery - from ._work_order import _WorkOrderQuery - - super().__init__(created_types, creation_path, client, result_list_cls, expression, None, connection_name) - if end_node_cls not in created_types: - self.end_node = end_node_cls( - created_types=created_types.copy(), - creation_path=self._creation_path, - client=client, - result_list_cls=result_list_cls, # type: ignore[type-var] - expression=dm.query.NodeResultSetExpression(), - ) - - self.space = StringFilter(self, ["node", "space"]) - self.external_id = StringFilter(self, ["node", "externalId"]) - self.end_time = TimestampFilter(self, self._view_id.as_property_ref("end_time")) - self.start_time = TimestampFilter(self, self._view_id.as_property_ref("start_time")) - self._filter_classes.extend( - [ - self.space, - self.external_id, - self.end_time, - self.start_time, - ] - ) diff --git a/examples/equipment_unit/data_classes/_unit_procedure.py b/examples/equipment_unit/data_classes/_unit_procedure.py deleted file mode 100644 index 975a40de0..000000000 --- a/examples/equipment_unit/data_classes/_unit_procedure.py +++ /dev/null @@ -1,518 +0,0 @@ -from __future__ import annotations - -import warnings -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, ClassVar, Literal, no_type_check, Optional, Union - -from cognite.client import data_modeling as dm, CogniteClient -from pydantic import Field -from pydantic import field_validator, model_validator - -from equipment_unit.data_classes._core import ( - DEFAULT_INSTANCE_SPACE, - DEFAULT_QUERY_LIMIT, - DataRecord, - DataRecordGraphQL, - DataRecordWrite, - DomainModel, - DomainModelWrite, - DomainModelWriteList, - DomainModelList, - DomainRelation, - DomainRelationWrite, - GraphQLCore, - ResourcesWrite, - T_DomainModelList, - as_direct_relation_reference, - as_instance_dict_id, - as_node_id, - as_pygen_node_id, - are_nodes_equal, - is_tuple_id, - select_best_node, - QueryCore, - NodeQueryCore, - StringFilter, -) - -if TYPE_CHECKING: - from equipment_unit.data_classes._start_end_time import ( - StartEndTime, - StartEndTimeList, - StartEndTimeGraphQL, - StartEndTimeWrite, - StartEndTimeWriteList, - ) - - -__all__ = [ - "UnitProcedure", - "UnitProcedureWrite", - "UnitProcedureApply", - "UnitProcedureList", - "UnitProcedureWriteList", - "UnitProcedureApplyList", - "UnitProcedureFields", - "UnitProcedureTextFields", - "UnitProcedureGraphQL", -] - - -UnitProcedureTextFields = Literal["external_id", "name", "type_"] -UnitProcedureFields = Literal["external_id", "name", "type_"] - -_UNITPROCEDURE_PROPERTIES_BY_FIELD = { - "external_id": "externalId", - "name": "name", - "type_": "type", -} - - -class UnitProcedureGraphQL(GraphQLCore): - """This represents the reading version of unit procedure, used - when data is retrieved from CDF using GraphQL. - - It is used when retrieving data from CDF using GraphQL. - - Args: - space: The space where the node is located. - external_id: The external id of the unit procedure. - data_record: The data record of the unit procedure node. - name: The name field. - type_: The type field. - work_orders: The work order field. - work_units: The work unit field. - """ - - view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "UnitProcedure", "a6e2fea1e1c664") - name: Optional[str] = None - type_: Optional[str] = Field(None, alias="type") - work_orders: Optional[list[StartEndTimeGraphQL]] = Field(default=None, repr=False) - work_units: Optional[list[StartEndTimeGraphQL]] = Field(default=None, repr=False) - - @model_validator(mode="before") - def parse_data_record(cls, values: Any) -> Any: - if not isinstance(values, dict): - return values - if "lastUpdatedTime" in values or "createdTime" in values: - values["dataRecord"] = DataRecordGraphQL( - created_time=values.pop("createdTime", None), - last_updated_time=values.pop("lastUpdatedTime", None), - ) - return values - - @field_validator("work_orders", "work_units", mode="before") - def parse_graphql(cls, value: Any) -> Any: - if not isinstance(value, dict): - return value - if "items" in value: - return value["items"] - return value - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_read(self) -> UnitProcedure: - """Convert this GraphQL format of unit procedure to the reading format.""" - if self.data_record is None: - raise ValueError("This object cannot be converted to a read format because it lacks a data record.") - return UnitProcedure( - space=self.space, - external_id=self.external_id, - data_record=DataRecord( - version=0, - last_updated_time=self.data_record.last_updated_time, - created_time=self.data_record.created_time, - ), - name=self.name, - type_=self.type_, - work_orders=[work_order.as_read() for work_order in self.work_orders or []], - work_units=[work_unit.as_read() for work_unit in self.work_units or []], - ) - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_write(self) -> UnitProcedureWrite: - """Convert this GraphQL format of unit procedure to the writing format.""" - return UnitProcedureWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=0), - name=self.name, - type_=self.type_, - work_orders=[work_order.as_write() for work_order in self.work_orders or []], - work_units=[work_unit.as_write() for work_unit in self.work_units or []], - ) - - -class UnitProcedure(DomainModel): - """This represents the reading version of unit procedure. - - It is used to when data is retrieved from CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the unit procedure. - data_record: The data record of the unit procedure node. - name: The name field. - type_: The type field. - work_orders: The work order field. - work_units: The work unit field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "UnitProcedure", "a6e2fea1e1c664") - - space: str = DEFAULT_INSTANCE_SPACE - node_type: Union[dm.DirectRelationReference, None] = None - name: Optional[str] = None - type_: Optional[str] = Field(None, alias="type") - work_orders: Optional[list[StartEndTime]] = Field(default=None, repr=False) - work_units: Optional[list[StartEndTime]] = Field(default=None, repr=False) - - def as_write(self) -> UnitProcedureWrite: - """Convert this read version of unit procedure to the writing version.""" - return UnitProcedureWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=self.data_record.version), - name=self.name, - type_=self.type_, - work_orders=[work_order.as_write() for work_order in self.work_orders or []], - work_units=[work_unit.as_write() for work_unit in self.work_units or []], - ) - - def as_apply(self) -> UnitProcedureWrite: - """Convert this read version of unit procedure to the writing version.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - @classmethod - def _update_connections( - cls, - instances: dict[dm.NodeId | str, UnitProcedure], # type: ignore[override] - nodes_by_id: dict[dm.NodeId | str, DomainModel], - edges_by_source_node: dict[dm.NodeId, list[dm.Edge | DomainRelation]], - ) -> None: - from ._start_end_time import StartEndTime - - for instance in instances.values(): - if edges := edges_by_source_node.get(instance.as_id()): - work_orders: list[StartEndTime] = [] - work_units: list[StartEndTime] = [] - for edge in edges: - value: DomainModel | DomainRelation | str | dm.NodeId - if isinstance(edge, DomainRelation): - value = edge - else: - other_end: dm.DirectRelationReference = ( - edge.end_node - if edge.start_node.space == instance.space - and edge.start_node.external_id == instance.external_id - else edge.start_node - ) - destination: dm.NodeId | str = ( - as_node_id(other_end) - if other_end.space != DEFAULT_INSTANCE_SPACE - else other_end.external_id - ) - if destination in nodes_by_id: - value = nodes_by_id[destination] - else: - value = destination - edge_type = edge.edge_type if isinstance(edge, DomainRelation) else edge.type - - if edge_type == dm.DirectRelationReference( - "IntegrationTestsImmutable", "UnitProcedure.work_order" - ) and isinstance(value, StartEndTime): - work_orders.append(value) - if end_node := nodes_by_id.get(as_pygen_node_id(value.end_node)): - value.end_node = end_node # type: ignore[assignment] - if edge_type == dm.DirectRelationReference( - "IntegrationTestsImmutable", "UnitProcedure.equipment_module" - ) and isinstance(value, StartEndTime): - work_units.append(value) - if end_node := nodes_by_id.get(as_pygen_node_id(value.end_node)): - value.end_node = end_node # type: ignore[assignment] - - instance.work_orders = work_orders - instance.work_units = work_units - - -class UnitProcedureWrite(DomainModelWrite): - """This represents the writing version of unit procedure. - - It is used to when data is sent to CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the unit procedure. - data_record: The data record of the unit procedure node. - name: The name field. - type_: The type field. - work_orders: The work order field. - work_units: The work unit field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "UnitProcedure", "a6e2fea1e1c664") - - space: str = DEFAULT_INSTANCE_SPACE - node_type: Union[dm.DirectRelationReference, dm.NodeId, tuple[str, str], None] = None - name: Optional[str] = None - type_: Optional[str] = Field(None, alias="type") - work_orders: Optional[list[StartEndTimeWrite]] = Field(default=None, repr=False) - work_units: Optional[list[StartEndTimeWrite]] = Field(default=None, repr=False) - - @field_validator("work_orders", "work_units", mode="before") - def as_node_id(cls, value: Any) -> Any: - if isinstance(value, dm.DirectRelationReference): - return dm.NodeId(value.space, value.external_id) - elif isinstance(value, tuple) and len(value) == 2 and all(isinstance(item, str) for item in value): - return dm.NodeId(value[0], value[1]) - elif isinstance(value, list): - return [cls.as_node_id(item) for item in value] - return value - - def _to_instances_write( - self, - cache: set[tuple[str, str]], - write_none: bool = False, - allow_version_increase: bool = False, - ) -> ResourcesWrite: - resources = ResourcesWrite() - if self.as_tuple_id() in cache: - return resources - - properties: dict[str, Any] = {} - - if self.name is not None or write_none: - properties["name"] = self.name - - if self.type_ is not None or write_none: - properties["type"] = self.type_ - - if properties: - this_node = dm.NodeApply( - space=self.space, - external_id=self.external_id, - existing_version=None if allow_version_increase else self.data_record.existing_version, - type=as_direct_relation_reference(self.node_type), - sources=[ - dm.NodeOrEdgeData( - source=self._view_id, - properties=properties, - ) - ], - ) - resources.nodes.append(this_node) - cache.add(self.as_tuple_id()) - - for work_order in self.work_orders or []: - if isinstance(work_order, DomainRelationWrite): - other_resources = work_order._to_instances_write( - cache, - self, - dm.DirectRelationReference("IntegrationTestsImmutable", "UnitProcedure.work_order"), - ) - resources.extend(other_resources) - - for work_unit in self.work_units or []: - if isinstance(work_unit, DomainRelationWrite): - other_resources = work_unit._to_instances_write( - cache, - self, - dm.DirectRelationReference("IntegrationTestsImmutable", "UnitProcedure.equipment_module"), - ) - resources.extend(other_resources) - - return resources - - -class UnitProcedureApply(UnitProcedureWrite): - def __new__(cls, *args, **kwargs) -> UnitProcedureApply: - warnings.warn( - "UnitProcedureApply is deprecated and will be removed in v1.0. Use UnitProcedureWrite instead." - "The motivation for this change is that Write is a more descriptive name for the writing version of the" - "UnitProcedure.", - UserWarning, - stacklevel=2, - ) - return super().__new__(cls) - - -class UnitProcedureList(DomainModelList[UnitProcedure]): - """List of unit procedures in the read version.""" - - _INSTANCE = UnitProcedure - - def as_write(self) -> UnitProcedureWriteList: - """Convert these read versions of unit procedure to the writing versions.""" - return UnitProcedureWriteList([node.as_write() for node in self.data]) - - def as_apply(self) -> UnitProcedureWriteList: - """Convert these read versions of primitive nullable to the writing versions.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - @property - def work_orders(self) -> StartEndTimeList: - from ._start_end_time import StartEndTime, StartEndTimeList - - return StartEndTimeList( - [item for items in self.data for item in items.work_orders or [] if isinstance(item, StartEndTime)] - ) - - @property - def work_units(self) -> StartEndTimeList: - from ._start_end_time import StartEndTime, StartEndTimeList - - return StartEndTimeList( - [item for items in self.data for item in items.work_units or [] if isinstance(item, StartEndTime)] - ) - - -class UnitProcedureWriteList(DomainModelWriteList[UnitProcedureWrite]): - """List of unit procedures in the writing version.""" - - _INSTANCE = UnitProcedureWrite - - @property - def work_orders(self) -> StartEndTimeWriteList: - from ._start_end_time import StartEndTimeWrite, StartEndTimeWriteList - - return StartEndTimeWriteList( - [item for items in self.data for item in items.work_orders or [] if isinstance(item, StartEndTimeWrite)] - ) - - @property - def work_units(self) -> StartEndTimeWriteList: - from ._start_end_time import StartEndTimeWrite, StartEndTimeWriteList - - return StartEndTimeWriteList( - [item for items in self.data for item in items.work_units or [] if isinstance(item, StartEndTimeWrite)] - ) - - -class UnitProcedureApplyList(UnitProcedureWriteList): ... - - -def _create_unit_procedure_filter( - view_id: dm.ViewId, - name: str | list[str] | None = None, - name_prefix: str | None = None, - type_: str | list[str] | None = None, - type_prefix: str | None = None, - external_id_prefix: str | None = None, - space: str | list[str] | None = None, - filter: dm.Filter | None = None, -) -> dm.Filter | None: - filters: list[dm.Filter] = [] - if isinstance(name, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("name"), value=name)) - if name and isinstance(name, list): - filters.append(dm.filters.In(view_id.as_property_ref("name"), values=name)) - if name_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("name"), value=name_prefix)) - if isinstance(type_, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("type"), value=type_)) - if type_ and isinstance(type_, list): - filters.append(dm.filters.In(view_id.as_property_ref("type"), values=type_)) - if type_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("type"), value=type_prefix)) - if external_id_prefix is not None: - filters.append(dm.filters.Prefix(["node", "externalId"], value=external_id_prefix)) - if isinstance(space, str): - filters.append(dm.filters.Equals(["node", "space"], value=space)) - if space and isinstance(space, list): - filters.append(dm.filters.In(["node", "space"], values=space)) - if filter: - filters.append(filter) - return dm.filters.And(*filters) if filters else None - - -class _UnitProcedureQuery(NodeQueryCore[T_DomainModelList, UnitProcedureList]): - _view_id = UnitProcedure._view_id - _result_cls = UnitProcedure - _result_list_cls_end = UnitProcedureList - - def __init__( - self, - created_types: set[type], - creation_path: list[QueryCore], - client: CogniteClient, - result_list_cls: type[T_DomainModelList], - expression: dm.query.ResultSetExpression | None = None, - connection_name: str | None = None, - connection_type: Literal["reverse-list"] | None = None, - reverse_expression: dm.query.ResultSetExpression | None = None, - ): - from ._equipment_module import _EquipmentModuleQuery - from ._start_end_time import _StartEndTimeQuery - from ._work_order import _WorkOrderQuery - - super().__init__( - created_types, - creation_path, - client, - result_list_cls, - expression, - dm.filters.HasData(views=[self._view_id]), - connection_name, - connection_type, - reverse_expression, - ) - - if _StartEndTimeQuery not in created_types: - self.work_orders = _StartEndTimeQuery( - created_types.copy(), - self._creation_path, - client, - result_list_cls, - _WorkOrderQuery, - dm.query.EdgeResultSetExpression( - direction="outwards", - chain_to="destination", - ), - connection_name="work_orders", - ) - - if _StartEndTimeQuery not in created_types: - self.work_units = _StartEndTimeQuery( - created_types.copy(), - self._creation_path, - client, - result_list_cls, - _EquipmentModuleQuery, - dm.query.EdgeResultSetExpression( - direction="outwards", - chain_to="destination", - ), - connection_name="work_units", - ) - - self.space = StringFilter(self, ["node", "space"]) - self.external_id = StringFilter(self, ["node", "externalId"]) - self.name = StringFilter(self, self._view_id.as_property_ref("name")) - self.type_ = StringFilter(self, self._view_id.as_property_ref("type")) - self._filter_classes.extend( - [ - self.space, - self.external_id, - self.name, - self.type_, - ] - ) - - def list_unit_procedure(self, limit: int = DEFAULT_QUERY_LIMIT) -> UnitProcedureList: - return self._list(limit=limit) - - -class UnitProcedureQuery(_UnitProcedureQuery[UnitProcedureList]): - def __init__(self, client: CogniteClient): - super().__init__(set(), [], client, UnitProcedureList) diff --git a/examples/equipment_unit/data_classes/_work_order.py b/examples/equipment_unit/data_classes/_work_order.py deleted file mode 100644 index 81b49d59c..000000000 --- a/examples/equipment_unit/data_classes/_work_order.py +++ /dev/null @@ -1,364 +0,0 @@ -from __future__ import annotations - -import warnings -from collections.abc import Sequence -from typing import Any, ClassVar, Literal, no_type_check, Optional, Union - -from cognite.client import data_modeling as dm, CogniteClient -from pydantic import Field -from pydantic import field_validator, model_validator - -from equipment_unit.data_classes._core import ( - DEFAULT_INSTANCE_SPACE, - DEFAULT_QUERY_LIMIT, - DataRecord, - DataRecordGraphQL, - DataRecordWrite, - DomainModel, - DomainModelWrite, - DomainModelWriteList, - DomainModelList, - DomainRelation, - DomainRelationWrite, - GraphQLCore, - ResourcesWrite, - T_DomainModelList, - as_direct_relation_reference, - as_instance_dict_id, - as_node_id, - as_pygen_node_id, - are_nodes_equal, - is_tuple_id, - select_best_node, - QueryCore, - NodeQueryCore, - StringFilter, -) - - -__all__ = [ - "WorkOrder", - "WorkOrderWrite", - "WorkOrderApply", - "WorkOrderList", - "WorkOrderWriteList", - "WorkOrderApplyList", - "WorkOrderFields", - "WorkOrderTextFields", - "WorkOrderGraphQL", -] - - -WorkOrderTextFields = Literal["external_id", "description", "performed_by", "type_"] -WorkOrderFields = Literal["external_id", "description", "performed_by", "type_"] - -_WORKORDER_PROPERTIES_BY_FIELD = { - "external_id": "externalId", - "description": "description", - "performed_by": "performedBy", - "type_": "type", -} - - -class WorkOrderGraphQL(GraphQLCore): - """This represents the reading version of work order, used - when data is retrieved from CDF using GraphQL. - - It is used when retrieving data from CDF using GraphQL. - - Args: - space: The space where the node is located. - external_id: The external id of the work order. - data_record: The data record of the work order node. - description: The description field. - performed_by: The performed by field. - type_: The type field. - """ - - view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "WorkOrder", "c5543fb2b1bc81") - description: Optional[str] = None - performed_by: Optional[str] = Field(None, alias="performedBy") - type_: Optional[str] = Field(None, alias="type") - - @model_validator(mode="before") - def parse_data_record(cls, values: Any) -> Any: - if not isinstance(values, dict): - return values - if "lastUpdatedTime" in values or "createdTime" in values: - values["dataRecord"] = DataRecordGraphQL( - created_time=values.pop("createdTime", None), - last_updated_time=values.pop("lastUpdatedTime", None), - ) - return values - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_read(self) -> WorkOrder: - """Convert this GraphQL format of work order to the reading format.""" - if self.data_record is None: - raise ValueError("This object cannot be converted to a read format because it lacks a data record.") - return WorkOrder( - space=self.space, - external_id=self.external_id, - data_record=DataRecord( - version=0, - last_updated_time=self.data_record.last_updated_time, - created_time=self.data_record.created_time, - ), - description=self.description, - performed_by=self.performed_by, - type_=self.type_, - ) - - # We do the ignore argument type as we let pydantic handle the type checking - @no_type_check - def as_write(self) -> WorkOrderWrite: - """Convert this GraphQL format of work order to the writing format.""" - return WorkOrderWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=0), - description=self.description, - performed_by=self.performed_by, - type_=self.type_, - ) - - -class WorkOrder(DomainModel): - """This represents the reading version of work order. - - It is used to when data is retrieved from CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the work order. - data_record: The data record of the work order node. - description: The description field. - performed_by: The performed by field. - type_: The type field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "WorkOrder", "c5543fb2b1bc81") - - space: str = DEFAULT_INSTANCE_SPACE - node_type: Union[dm.DirectRelationReference, None] = None - description: Optional[str] = None - performed_by: Optional[str] = Field(None, alias="performedBy") - type_: Optional[str] = Field(None, alias="type") - - def as_write(self) -> WorkOrderWrite: - """Convert this read version of work order to the writing version.""" - return WorkOrderWrite( - space=self.space, - external_id=self.external_id, - data_record=DataRecordWrite(existing_version=self.data_record.version), - description=self.description, - performed_by=self.performed_by, - type_=self.type_, - ) - - def as_apply(self) -> WorkOrderWrite: - """Convert this read version of work order to the writing version.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - -class WorkOrderWrite(DomainModelWrite): - """This represents the writing version of work order. - - It is used to when data is sent to CDF. - - Args: - space: The space where the node is located. - external_id: The external id of the work order. - data_record: The data record of the work order node. - description: The description field. - performed_by: The performed by field. - type_: The type field. - """ - - _view_id: ClassVar[dm.ViewId] = dm.ViewId("IntegrationTestsImmutable", "WorkOrder", "c5543fb2b1bc81") - - space: str = DEFAULT_INSTANCE_SPACE - node_type: Union[dm.DirectRelationReference, dm.NodeId, tuple[str, str], None] = None - description: Optional[str] = None - performed_by: Optional[str] = Field(None, alias="performedBy") - type_: Optional[str] = Field(None, alias="type") - - def _to_instances_write( - self, - cache: set[tuple[str, str]], - write_none: bool = False, - allow_version_increase: bool = False, - ) -> ResourcesWrite: - resources = ResourcesWrite() - if self.as_tuple_id() in cache: - return resources - - properties: dict[str, Any] = {} - - if self.description is not None or write_none: - properties["description"] = self.description - - if self.performed_by is not None or write_none: - properties["performedBy"] = self.performed_by - - if self.type_ is not None or write_none: - properties["type"] = self.type_ - - if properties: - this_node = dm.NodeApply( - space=self.space, - external_id=self.external_id, - existing_version=None if allow_version_increase else self.data_record.existing_version, - type=as_direct_relation_reference(self.node_type), - sources=[ - dm.NodeOrEdgeData( - source=self._view_id, - properties=properties, - ) - ], - ) - resources.nodes.append(this_node) - cache.add(self.as_tuple_id()) - - return resources - - -class WorkOrderApply(WorkOrderWrite): - def __new__(cls, *args, **kwargs) -> WorkOrderApply: - warnings.warn( - "WorkOrderApply is deprecated and will be removed in v1.0. Use WorkOrderWrite instead." - "The motivation for this change is that Write is a more descriptive name for the writing version of the" - "WorkOrder.", - UserWarning, - stacklevel=2, - ) - return super().__new__(cls) - - -class WorkOrderList(DomainModelList[WorkOrder]): - """List of work orders in the read version.""" - - _INSTANCE = WorkOrder - - def as_write(self) -> WorkOrderWriteList: - """Convert these read versions of work order to the writing versions.""" - return WorkOrderWriteList([node.as_write() for node in self.data]) - - def as_apply(self) -> WorkOrderWriteList: - """Convert these read versions of primitive nullable to the writing versions.""" - warnings.warn( - "as_apply is deprecated and will be removed in v1.0. Use as_write instead.", - UserWarning, - stacklevel=2, - ) - return self.as_write() - - -class WorkOrderWriteList(DomainModelWriteList[WorkOrderWrite]): - """List of work orders in the writing version.""" - - _INSTANCE = WorkOrderWrite - - -class WorkOrderApplyList(WorkOrderWriteList): ... - - -def _create_work_order_filter( - view_id: dm.ViewId, - description: str | list[str] | None = None, - description_prefix: str | None = None, - performed_by: str | list[str] | None = None, - performed_by_prefix: str | None = None, - type_: str | list[str] | None = None, - type_prefix: str | None = None, - external_id_prefix: str | None = None, - space: str | list[str] | None = None, - filter: dm.Filter | None = None, -) -> dm.Filter | None: - filters: list[dm.Filter] = [] - if isinstance(description, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("description"), value=description)) - if description and isinstance(description, list): - filters.append(dm.filters.In(view_id.as_property_ref("description"), values=description)) - if description_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("description"), value=description_prefix)) - if isinstance(performed_by, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("performedBy"), value=performed_by)) - if performed_by and isinstance(performed_by, list): - filters.append(dm.filters.In(view_id.as_property_ref("performedBy"), values=performed_by)) - if performed_by_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("performedBy"), value=performed_by_prefix)) - if isinstance(type_, str): - filters.append(dm.filters.Equals(view_id.as_property_ref("type"), value=type_)) - if type_ and isinstance(type_, list): - filters.append(dm.filters.In(view_id.as_property_ref("type"), values=type_)) - if type_prefix is not None: - filters.append(dm.filters.Prefix(view_id.as_property_ref("type"), value=type_prefix)) - if external_id_prefix is not None: - filters.append(dm.filters.Prefix(["node", "externalId"], value=external_id_prefix)) - if isinstance(space, str): - filters.append(dm.filters.Equals(["node", "space"], value=space)) - if space and isinstance(space, list): - filters.append(dm.filters.In(["node", "space"], values=space)) - if filter: - filters.append(filter) - return dm.filters.And(*filters) if filters else None - - -class _WorkOrderQuery(NodeQueryCore[T_DomainModelList, WorkOrderList]): - _view_id = WorkOrder._view_id - _result_cls = WorkOrder - _result_list_cls_end = WorkOrderList - - def __init__( - self, - created_types: set[type], - creation_path: list[QueryCore], - client: CogniteClient, - result_list_cls: type[T_DomainModelList], - expression: dm.query.ResultSetExpression | None = None, - connection_name: str | None = None, - connection_type: Literal["reverse-list"] | None = None, - reverse_expression: dm.query.ResultSetExpression | None = None, - ): - - super().__init__( - created_types, - creation_path, - client, - result_list_cls, - expression, - dm.filters.HasData(views=[self._view_id]), - connection_name, - connection_type, - reverse_expression, - ) - - self.space = StringFilter(self, ["node", "space"]) - self.external_id = StringFilter(self, ["node", "externalId"]) - self.description = StringFilter(self, self._view_id.as_property_ref("description")) - self.performed_by = StringFilter(self, self._view_id.as_property_ref("performedBy")) - self.type_ = StringFilter(self, self._view_id.as_property_ref("type")) - self._filter_classes.extend( - [ - self.space, - self.external_id, - self.description, - self.performed_by, - self.type_, - ] - ) - - def list_work_order(self, limit: int = DEFAULT_QUERY_LIMIT) -> WorkOrderList: - return self._list(limit=limit) - - -class WorkOrderQuery(_WorkOrderQuery[WorkOrderList]): - def __init__(self, client: CogniteClient): - super().__init__(set(), [], client, WorkOrderList) diff --git a/examples/omni/_api/cdf_external_references.py b/examples/omni/_api/cdf_external_references.py index 74f1cd784..96815b950 100644 --- a/examples/omni/_api/cdf_external_references.py +++ b/examples/omni/_api/cdf_external_references.py @@ -76,6 +76,11 @@ def __call__( A query API for cdf external references. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_cdf_external_reference_filter( self._view_id, diff --git a/examples/omni/_api/cdf_external_references_listed.py b/examples/omni/_api/cdf_external_references_listed.py index 66a6b919a..4409305bc 100644 --- a/examples/omni/_api/cdf_external_references_listed.py +++ b/examples/omni/_api/cdf_external_references_listed.py @@ -79,6 +79,11 @@ def __call__( A query API for cdf external references listeds. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_cdf_external_references_listed_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_a.py b/examples/omni/_api/connection_item_a.py index cd6ebe035..06413eba8 100644 --- a/examples/omni/_api/connection_item_a.py +++ b/examples/omni/_api/connection_item_a.py @@ -96,6 +96,11 @@ def __call__( A query API for connection item as. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_a_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_b.py b/examples/omni/_api/connection_item_b.py index 6d4f71bf1..c927d6ebd 100644 --- a/examples/omni/_api/connection_item_b.py +++ b/examples/omni/_api/connection_item_b.py @@ -79,6 +79,11 @@ def __call__( A query API for connection item bs. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_b_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_c_node.py b/examples/omni/_api/connection_item_c_node.py index 2be295fd0..d3728a57e 100644 --- a/examples/omni/_api/connection_item_c_node.py +++ b/examples/omni/_api/connection_item_c_node.py @@ -77,6 +77,11 @@ def __call__( A query API for connection item c nodes. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_c_node_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_d.py b/examples/omni/_api/connection_item_d.py index cdaa0fe93..b7ccc87d0 100644 --- a/examples/omni/_api/connection_item_d.py +++ b/examples/omni/_api/connection_item_d.py @@ -95,6 +95,11 @@ def __call__( A query API for connection item ds. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_d_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_e.py b/examples/omni/_api/connection_item_e.py index 7ddd30d27..eb34626da 100644 --- a/examples/omni/_api/connection_item_e.py +++ b/examples/omni/_api/connection_item_e.py @@ -93,6 +93,11 @@ def __call__( A query API for connection item es. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_e_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_f.py b/examples/omni/_api/connection_item_f.py index af98622d5..190e563f3 100644 --- a/examples/omni/_api/connection_item_f.py +++ b/examples/omni/_api/connection_item_f.py @@ -94,6 +94,11 @@ def __call__( A query API for connection item fs. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_f_filter( self._view_id, diff --git a/examples/omni/_api/connection_item_g.py b/examples/omni/_api/connection_item_g.py index 11fd8f519..d453e6ac5 100644 --- a/examples/omni/_api/connection_item_g.py +++ b/examples/omni/_api/connection_item_g.py @@ -81,6 +81,11 @@ def __call__( A query API for connection item gs. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_g_filter( self._view_id, diff --git a/examples/omni/_api/dependent_on_non_writable.py b/examples/omni/_api/dependent_on_non_writable.py index 2780d3a23..eb04fb667 100644 --- a/examples/omni/_api/dependent_on_non_writable.py +++ b/examples/omni/_api/dependent_on_non_writable.py @@ -81,6 +81,11 @@ def __call__( A query API for dependent on non writables. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_dependent_on_non_writable_filter( self._view_id, diff --git a/examples/omni/_api/empty.py b/examples/omni/_api/empty.py index 40a20f568..74a91cb06 100644 --- a/examples/omni/_api/empty.py +++ b/examples/omni/_api/empty.py @@ -100,6 +100,11 @@ def __call__( A query API for empties. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_empty_filter( self._view_id, diff --git a/examples/omni/_api/implementation_1.py b/examples/omni/_api/implementation_1.py index a0bf0989b..fd26d000a 100644 --- a/examples/omni/_api/implementation_1.py +++ b/examples/omni/_api/implementation_1.py @@ -85,6 +85,11 @@ def __call__( A query API for implementation 1. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_implementation_1_filter( self._view_id, diff --git a/examples/omni/_api/implementation_1_non_writeable.py b/examples/omni/_api/implementation_1_non_writeable.py index 32e575f68..be8917df4 100644 --- a/examples/omni/_api/implementation_1_non_writeable.py +++ b/examples/omni/_api/implementation_1_non_writeable.py @@ -78,6 +78,11 @@ def __call__( A query API for implementation 1 non writeables. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_implementation_1_non_writeable_filter( self._view_id, diff --git a/examples/omni/_api/implementation_2.py b/examples/omni/_api/implementation_2.py index 73916b376..1e0d56905 100644 --- a/examples/omni/_api/implementation_2.py +++ b/examples/omni/_api/implementation_2.py @@ -77,6 +77,11 @@ def __call__( A query API for implementation 2. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_implementation_2_filter( self._view_id, diff --git a/examples/omni/_api/main_interface.py b/examples/omni/_api/main_interface.py index d4499ea6c..9126218dd 100644 --- a/examples/omni/_api/main_interface.py +++ b/examples/omni/_api/main_interface.py @@ -79,6 +79,11 @@ def __call__( A query API for main interfaces. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_main_interface_filter( self._view_id, diff --git a/examples/omni/_api/primitive_nullable.py b/examples/omni/_api/primitive_nullable.py index 1dd98cc39..3672a0341 100644 --- a/examples/omni/_api/primitive_nullable.py +++ b/examples/omni/_api/primitive_nullable.py @@ -102,6 +102,11 @@ def __call__( A query API for primitive nullables. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_primitive_nullable_filter( self._view_id, diff --git a/examples/omni/_api/primitive_nullable_listed.py b/examples/omni/_api/primitive_nullable_listed.py index 9c90ef8cd..090084a36 100644 --- a/examples/omni/_api/primitive_nullable_listed.py +++ b/examples/omni/_api/primitive_nullable_listed.py @@ -77,6 +77,11 @@ def __call__( A query API for primitive nullable listeds. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_primitive_nullable_listed_filter( self._view_id, diff --git a/examples/omni/_api/primitive_required.py b/examples/omni/_api/primitive_required.py index d531b4e0d..f8e6f7ef6 100644 --- a/examples/omni/_api/primitive_required.py +++ b/examples/omni/_api/primitive_required.py @@ -102,6 +102,11 @@ def __call__( A query API for primitive requireds. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_primitive_required_filter( self._view_id, diff --git a/examples/omni/_api/primitive_required_listed.py b/examples/omni/_api/primitive_required_listed.py index 9943bac0c..e6791773f 100644 --- a/examples/omni/_api/primitive_required_listed.py +++ b/examples/omni/_api/primitive_required_listed.py @@ -77,6 +77,11 @@ def __call__( A query API for primitive required listeds. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_primitive_required_listed_filter( self._view_id, diff --git a/examples/omni/_api/primitive_with_defaults.py b/examples/omni/_api/primitive_with_defaults.py index 18d8c6144..dfc82409c 100644 --- a/examples/omni/_api/primitive_with_defaults.py +++ b/examples/omni/_api/primitive_with_defaults.py @@ -87,6 +87,11 @@ def __call__( A query API for primitive with defaults. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_primitive_with_default_filter( self._view_id, diff --git a/examples/omni/_api/sub_interface.py b/examples/omni/_api/sub_interface.py index abe498e7e..bbae7417c 100644 --- a/examples/omni/_api/sub_interface.py +++ b/examples/omni/_api/sub_interface.py @@ -85,6 +85,11 @@ def __call__( A query API for sub interfaces. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_sub_interface_filter( self._view_id, diff --git a/examples/omni_multi/_api/implementation_1_v_1.py b/examples/omni_multi/_api/implementation_1_v_1.py index b4b3e0a23..3d955f55c 100644 --- a/examples/omni_multi/_api/implementation_1_v_1.py +++ b/examples/omni_multi/_api/implementation_1_v_1.py @@ -82,6 +82,11 @@ def __call__( A query API for implementation 1 v 1. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_implementation_1_v_1_filter( self._view_id, diff --git a/examples/omni_multi/_api/implementation_1_v_2.py b/examples/omni_multi/_api/implementation_1_v_2.py index 3305bb132..f17e61c33 100644 --- a/examples/omni_multi/_api/implementation_1_v_2.py +++ b/examples/omni_multi/_api/implementation_1_v_2.py @@ -82,6 +82,11 @@ def __call__( A query API for implementation 1 v 2. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_implementation_1_v_2_filter( self._view_id, diff --git a/examples/omni_multi/_api/main_interface.py b/examples/omni_multi/_api/main_interface.py index 8fff3de7a..efe16645c 100644 --- a/examples/omni_multi/_api/main_interface.py +++ b/examples/omni_multi/_api/main_interface.py @@ -76,6 +76,11 @@ def __call__( A query API for main interfaces. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_main_interface_filter( self._view_id, diff --git a/examples/omni_multi/_api/sub_interface.py b/examples/omni_multi/_api/sub_interface.py index 4a4c74cb9..0a553277d 100644 --- a/examples/omni_multi/_api/sub_interface.py +++ b/examples/omni_multi/_api/sub_interface.py @@ -80,6 +80,11 @@ def __call__( A query API for sub interfaces. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_sub_interface_filter( self._view_id, diff --git a/examples/omni_sub/_api/connection_item_a.py b/examples/omni_sub/_api/connection_item_a.py index 053bc7f36..b3b8a8e2d 100644 --- a/examples/omni_sub/_api/connection_item_a.py +++ b/examples/omni_sub/_api/connection_item_a.py @@ -93,6 +93,11 @@ def __call__( A query API for connection item as. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_a_filter( self._view_id, diff --git a/examples/omni_sub/_api/connection_item_b.py b/examples/omni_sub/_api/connection_item_b.py index 9610f3dce..fe513cc0a 100644 --- a/examples/omni_sub/_api/connection_item_b.py +++ b/examples/omni_sub/_api/connection_item_b.py @@ -78,6 +78,11 @@ def __call__( A query API for connection item bs. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_b_filter( self._view_id, diff --git a/examples/omni_sub/_api/connection_item_c_node.py b/examples/omni_sub/_api/connection_item_c_node.py index bc76f857b..851c29e12 100644 --- a/examples/omni_sub/_api/connection_item_c_node.py +++ b/examples/omni_sub/_api/connection_item_c_node.py @@ -76,6 +76,11 @@ def __call__( A query API for connection item c nodes. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_connection_item_c_node_filter( self._view_id, diff --git a/examples/wind_turbine/_api/blade.py b/examples/wind_turbine/_api/blade.py index d8da3457b..2749a6cff 100644 --- a/examples/wind_turbine/_api/blade.py +++ b/examples/wind_turbine/_api/blade.py @@ -76,6 +76,11 @@ def __call__( A query API for blades. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_blade_filter( self._view_id, diff --git a/examples/wind_turbine/_api/data_sheet.py b/examples/wind_turbine/_api/data_sheet.py index 312c163f0..313e759ac 100644 --- a/examples/wind_turbine/_api/data_sheet.py +++ b/examples/wind_turbine/_api/data_sheet.py @@ -92,6 +92,11 @@ def __call__( A query API for data sheets. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_data_sheet_filter( self._view_id, diff --git a/examples/wind_turbine/_api/gearbox.py b/examples/wind_turbine/_api/gearbox.py index 5c4bd386e..33e24a3ce 100644 --- a/examples/wind_turbine/_api/gearbox.py +++ b/examples/wind_turbine/_api/gearbox.py @@ -97,6 +97,11 @@ def __call__( A query API for gearboxes. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_gearbox_filter( self._view_id, diff --git a/examples/wind_turbine/_api/generating_unit.py b/examples/wind_turbine/_api/generating_unit.py index 63cb4bd80..30de97aed 100644 --- a/examples/wind_turbine/_api/generating_unit.py +++ b/examples/wind_turbine/_api/generating_unit.py @@ -87,6 +87,11 @@ def __call__( A query API for generating units. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_generating_unit_filter( self._view_id, diff --git a/examples/wind_turbine/_api/generator.py b/examples/wind_turbine/_api/generator.py index a9af83a42..933b96cce 100644 --- a/examples/wind_turbine/_api/generator.py +++ b/examples/wind_turbine/_api/generator.py @@ -88,6 +88,11 @@ def __call__( A query API for generators. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_generator_filter( self._view_id, diff --git a/examples/wind_turbine/_api/high_speed_shaft.py b/examples/wind_turbine/_api/high_speed_shaft.py index 0514da60e..0203f8860 100644 --- a/examples/wind_turbine/_api/high_speed_shaft.py +++ b/examples/wind_turbine/_api/high_speed_shaft.py @@ -97,6 +97,11 @@ def __call__( A query API for high speed shafts. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_high_speed_shaft_filter( self._view_id, diff --git a/examples/wind_turbine/_api/main_shaft.py b/examples/wind_turbine/_api/main_shaft.py index 28e7400f5..a80b52a79 100644 --- a/examples/wind_turbine/_api/main_shaft.py +++ b/examples/wind_turbine/_api/main_shaft.py @@ -115,6 +115,11 @@ def __call__( A query API for main shafts. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_main_shaft_filter( self._view_id, diff --git a/examples/wind_turbine/_api/metmast.py b/examples/wind_turbine/_api/metmast.py index 6399683d7..5d8a4f276 100644 --- a/examples/wind_turbine/_api/metmast.py +++ b/examples/wind_turbine/_api/metmast.py @@ -87,6 +87,11 @@ def __call__( A query API for metmasts. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_metmast_filter( self._view_id, diff --git a/examples/wind_turbine/_api/nacelle.py b/examples/wind_turbine/_api/nacelle.py index 98847c22a..9e612817b 100644 --- a/examples/wind_turbine/_api/nacelle.py +++ b/examples/wind_turbine/_api/nacelle.py @@ -165,6 +165,11 @@ def __call__( A query API for nacelles. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_nacelle_filter( self._view_id, diff --git a/examples/wind_turbine/_api/power_inverter.py b/examples/wind_turbine/_api/power_inverter.py index 304f0d164..519660ada 100644 --- a/examples/wind_turbine/_api/power_inverter.py +++ b/examples/wind_turbine/_api/power_inverter.py @@ -97,6 +97,11 @@ def __call__( A query API for power inverters. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_power_inverter_filter( self._view_id, diff --git a/examples/wind_turbine/_api/rotor.py b/examples/wind_turbine/_api/rotor.py index f022f3ba5..2750adf74 100644 --- a/examples/wind_turbine/_api/rotor.py +++ b/examples/wind_turbine/_api/rotor.py @@ -88,6 +88,11 @@ def __call__( A query API for rotors. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_rotor_filter( self._view_id, diff --git a/examples/wind_turbine/_api/sensor_position.py b/examples/wind_turbine/_api/sensor_position.py index deed0d5a4..553455d08 100644 --- a/examples/wind_turbine/_api/sensor_position.py +++ b/examples/wind_turbine/_api/sensor_position.py @@ -156,6 +156,11 @@ def __call__( A query API for sensor positions. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_sensor_position_filter( self._view_id, diff --git a/examples/wind_turbine/_api/sensor_time_series.py b/examples/wind_turbine/_api/sensor_time_series.py index 94dd5f381..686a2554d 100644 --- a/examples/wind_turbine/_api/sensor_time_series.py +++ b/examples/wind_turbine/_api/sensor_time_series.py @@ -93,6 +93,11 @@ def __call__( A query API for sensor time series. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_sensor_time_series_filter( self._view_id, diff --git a/examples/wind_turbine/_api/solar_panel.py b/examples/wind_turbine/_api/solar_panel.py index feccc928b..e4d042934 100644 --- a/examples/wind_turbine/_api/solar_panel.py +++ b/examples/wind_turbine/_api/solar_panel.py @@ -100,6 +100,11 @@ def __call__( A query API for solar panels. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_solar_panel_filter( self._view_id, diff --git a/examples/wind_turbine/_api/wind_turbine.py b/examples/wind_turbine/_api/wind_turbine.py index f97b17364..5860be3b8 100644 --- a/examples/wind_turbine/_api/wind_turbine.py +++ b/examples/wind_turbine/_api/wind_turbine.py @@ -133,6 +133,11 @@ def __call__( A query API for wind turbines. """ + warnings.warn( + "This method is deprecated and will soon be removed. " "Use the .select() method instead.", + UserWarning, + stacklevel=2, + ) has_data = dm.filters.HasData(views=[self._view_id]) filter_ = _create_wind_turbine_filter( self._view_id, diff --git a/tests/constants.py b/tests/constants.py index 03a020fe5..9ee80c804 100644 --- a/tests/constants.py +++ b/tests/constants.py @@ -130,14 +130,6 @@ def load_read_nodes(self, data_model_id: dm.DataModelId) -> dm.NodeList: instance_space=None, ) -EQUIPMENT_UNIT_SDK = ExampleSDK( - data_model_ids=[DataModelId("IntegrationTestsImmutable", "EquipmentUnit", "2")], - _top_level_package="equipment_unit", - client_name="EquipmentUnitClient", - generate_sdk=True, - instance_space="IntegrationTestsImmutable", -) - CORE_SDK = ExampleSDK( data_model_ids=[DataModelId("cdf_cdm", "CogniteCore", "v1")], _top_level_package="cognite_core", @@ -163,27 +155,6 @@ def load_read_nodes(self, data_model_id: dm.DataModelId) -> dm.NodeList: ) -class EquipmentSDKFiles: - client_dir = EQUIPMENT_UNIT_SDK.client_dir - client = client_dir / "_api_client.py" - data_classes = client_dir / "data_classes" - core_data = data_classes / "_core.py" - start_end_time_data = data_classes / "_start_end_time.py" - unit_procedure_data = data_classes / "_unit_procedure.py" - equipment_module_data = data_classes / "_equipment_module.py" - - api = client_dir / "_api" - equipment_api = api / "equipment_module.py" - equipment_module_sensor_value_api = api / "equipment_module_sensor_value.py" - - unit_procedure_api = api / "unit_procedure.py" - unit_procedure_query = api / "unit_procedure_query.py" - unit_procedure_work_units = api / "unit_procedure_work_units.py" - core_api = api / "_core.py" - - data_init = data_classes / "__init__.py" - - class OmniFiles: client_dir = OMNI_SDK.client_dir client = client_dir / "_api_client.py" diff --git a/tests/data/models/EquipmentUnit/model.yaml b/tests/data/models/EquipmentUnit/model.yaml deleted file mode 100644 index 20345003d..000000000 --- a/tests/data/models/EquipmentUnit/model.yaml +++ /dev/null @@ -1,262 +0,0 @@ -space: IntegrationTestsImmutable -externalId: EquipmentUnit -name: EquipmentUnit -description: '' -version: '2' -views: -- space: IntegrationTestsImmutable - externalId: EquipmentModule - name: EquipmentModule - implements: [] - version: b1cd4bf14a7a33 - writable: true - usedFor: node - isGlobal: false - properties: - description: - container: - space: IntegrationTestsImmutable - externalId: EquipmentModule - containerPropertyIdentifier: description - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: description - description: null - name: - container: - space: IntegrationTestsImmutable - externalId: EquipmentModule - containerPropertyIdentifier: name - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: name - description: null - sensor_value: - container: - space: IntegrationTestsImmutable - externalId: EquipmentModule - containerPropertyIdentifier: sensor_value - type: - list: false - type: timeseries - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: sensor_value - description: null - type: - container: - space: IntegrationTestsImmutable - externalId: EquipmentModule - containerPropertyIdentifier: type - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: type - description: null - lastUpdatedTime: 1703424931271 - createdTime: 1703424931271 -- space: IntegrationTestsImmutable - externalId: StartEndTime - name: StartEndTime - implements: [] - version: d416e0ed98186b - writable: true - usedFor: edge - isGlobal: false - properties: - end_time: - container: - space: IntegrationTestsImmutable - externalId: StartEndTime - containerPropertyIdentifier: end_time - type: - list: false - type: timestamp - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: end_time - description: null - start_time: - container: - space: IntegrationTestsImmutable - externalId: StartEndTime - containerPropertyIdentifier: start_time - type: - list: false - type: timestamp - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: start_time - description: null - lastUpdatedTime: 1703424931271 - createdTime: 1703424931271 -- space: IntegrationTestsImmutable - externalId: UnitProcedure - name: UnitProcedure - implements: [] - version: a6e2fea1e1c664 - writable: true - usedFor: node - isGlobal: false - properties: - name: - container: - space: IntegrationTestsImmutable - externalId: UnitProcedure - containerPropertyIdentifier: name - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: name - description: null - type: - container: - space: IntegrationTestsImmutable - externalId: UnitProcedure - containerPropertyIdentifier: type - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: type - description: null - work_orders: - type: - space: IntegrationTestsImmutable - externalId: UnitProcedure.work_order - source: - space: IntegrationTestsImmutable - externalId: WorkOrder - version: c5543fb2b1bc81 - type: view - name: work_orders - description: null - edgeSource: - space: IntegrationTestsImmutable - externalId: StartEndTime - version: d416e0ed98186b - type: view - direction: outwards - connectionType: multi_edge_connection - work_units: - type: - space: IntegrationTestsImmutable - externalId: UnitProcedure.equipment_module - source: - space: IntegrationTestsImmutable - externalId: EquipmentModule - version: b1cd4bf14a7a33 - type: view - name: work_units - description: null - edgeSource: - space: IntegrationTestsImmutable - externalId: StartEndTime - version: d416e0ed98186b - type: view - direction: outwards - connectionType: multi_edge_connection - lastUpdatedTime: 1703424931271 - createdTime: 1703424931271 -- space: IntegrationTestsImmutable - externalId: WorkOrder - name: WorkOrder - implements: [] - version: c5543fb2b1bc81 - writable: true - usedFor: node - isGlobal: false - properties: - description: - container: - space: IntegrationTestsImmutable - externalId: WorkOrder - containerPropertyIdentifier: description - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: description - description: null - performedBy: - container: - space: IntegrationTestsImmutable - externalId: WorkOrder - containerPropertyIdentifier: performedBy - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: performedBy - description: null - type: - container: - space: IntegrationTestsImmutable - externalId: WorkOrder - containerPropertyIdentifier: type - type: - list: false - collation: ucs_basic - type: text - nullable: true - immutable: false - autoIncrement: false - source: null - defaultValue: null - name: type - description: null - lastUpdatedTime: 1703424931271 - createdTime: 1703424931271 -isGlobal: false -lastUpdatedTime: 1703485511412 -createdTime: 1703424932659 diff --git a/tests/test_integration/test_list_edges.py b/tests/test_integration/test_list_edges.py new file mode 100644 index 000000000..553f32eb6 --- /dev/null +++ b/tests/test_integration/test_list_edges.py @@ -0,0 +1,35 @@ +from wind_turbine import WindTurbineClient +from wind_turbine import data_classes as data_cls + + +def test_edges_to_pandas(turbine_client: WindTurbineClient) -> None: + df = turbine_client.wind_turbine.metmast_edge.list(max_distance=1000, limit=3).to_pandas() + assert len(df) <= 3 + assert not df.empty + assert "distance" in df.columns + assert sum(df["distance"] <= 1000) == len(df) + + +def test_upsert_with_edge(turbine_client: WindTurbineClient) -> None: + new_turbine = data_cls.WindTurbineWrite( + external_id="doctriono_b", + name="A new Wind Turbine", + capacity=8.0, + metmast=[ + data_cls.DistanceWrite( + distance=500.0, + end_node=data_cls.MetmastWrite( + external_id="doctrino_weather", + position=42.0, + ), + ) + ], + ) + + try: + created = turbine_client.upsert(new_turbine) + + assert len(created.nodes) == 2 + assert len(created.edges) == 1 + finally: + turbine_client.delete(new_turbine) diff --git a/tests/test_integration/test_sdks/__init__.py b/tests/test_integration/test_sdks/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/test_integration/test_sdks/test_equipment_sdk/__init__.py b/tests/test_integration/test_sdks/test_equipment_sdk/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/test_integration/test_sdks/test_equipment_sdk/conftest.py b/tests/test_integration/test_sdks/test_equipment_sdk/conftest.py deleted file mode 100644 index f4c097741..000000000 --- a/tests/test_integration/test_sdks/test_equipment_sdk/conftest.py +++ /dev/null @@ -1,7 +0,0 @@ -import pytest -from equipment_unit import EquipmentUnitClient - - -@pytest.fixture() -def workorder(client_config) -> EquipmentUnitClient: - return EquipmentUnitClient.azure_project(**client_config) diff --git a/tests/test_integration/test_sdks/test_equipment_sdk/test_unit_procedure.py b/tests/test_integration/test_sdks/test_equipment_sdk/test_unit_procedure.py deleted file mode 100644 index cafbeb676..000000000 --- a/tests/test_integration/test_sdks/test_equipment_sdk/test_unit_procedure.py +++ /dev/null @@ -1,104 +0,0 @@ -import pytest -from cognite.client import CogniteClient -from cognite.client.data_classes import TimeSeriesWrite -from equipment_unit.data_classes import ( - EquipmentModule, - EquipmentModuleWrite, - StartEndTime, - StartEndTimeList, - StartEndTimeWrite, - UnitProcedureList, - UnitProcedureWrite, -) - -from examples.equipment_unit import EquipmentUnitClient - - -@pytest.fixture -def start_end_time_edges(workorder: EquipmentUnitClient) -> StartEndTimeList: - edges = workorder.unit_procedure.work_units_edge.list(limit=-1) - assert len(edges) > 2, "There should be at least three edge in the list" - assert isinstance(edges, StartEndTimeList) - return edges - - -@pytest.fixture -def unit_procedure_list(workorder: EquipmentUnitClient) -> UnitProcedureList: - nodes = workorder.unit_procedure.list(limit=5) - assert len(nodes) > 2, "There should be at least three node in the list" - assert isinstance(nodes, UnitProcedureList) - return nodes - - -def test_edges_to_pandas(start_end_time_edges: StartEndTimeList) -> None: - df = start_end_time_edges.to_pandas() - assert len(df) == len(start_end_time_edges) - - -def test_unit_procedure_list_to_pandas(unit_procedure_list: UnitProcedureList) -> None: - df = unit_procedure_list.to_pandas() - assert len(df) == len(unit_procedure_list) - - -def test_single_unit_procedure_to_pandas(unit_procedure_list: UnitProcedureList) -> None: - procedure = unit_procedure_list[0] - - series = procedure.to_pandas() - assert len(series) == 8, "We only have the four properties + the external id, space, node type, and data_record" - - -def test_filter_start_end_time_edges(start_end_time_edges: StartEndTimeList, workorder: EquipmentUnitClient) -> None: - sorted_by_start_time = sorted(start_end_time_edges, key=lambda x: x.start_time) - - filtered = workorder.unit_procedure.work_units_edge.list( - min_start_time=sorted_by_start_time[1].start_time, limit=-1 - ) - - assert len(filtered) == len(sorted_by_start_time) - 1 - - -def test_filter_unit_procedure_through_edge(workorder: EquipmentUnitClient) -> None: - unit_procedures = workorder.unit_procedure(type_="red", limit=3).work_units(limit=5).query() - - assert 1 <= len(unit_procedures) <= 3 - assert all(procedure.type_ == "red" for procedure in unit_procedures) - for unit_procedure in unit_procedures: - for work_unit in unit_procedure.work_units: - assert isinstance(work_unit, StartEndTime) - assert isinstance(work_unit.end_node, EquipmentModule) - - -def test_apply_unit_procedure_with_edge(workorder: EquipmentUnitClient, cognite_client: CogniteClient) -> None: - new_procedure = UnitProcedureWrite( - external_id="procedure:new_procedure", - name="New procedure", - type_="New type", - work_units=[ - StartEndTimeWrite( - start_time="2021-01-01T00:00:00Z", - end_time="2021-01-01T00:00:00Z", - end_node=EquipmentModuleWrite( - external_id="module:new_module", - name="New module", - equipment_module_type="New type", - sensor_value=TimeSeriesWrite( - external_id="timeseries:123", - is_step=True, - description="This is a test timeseries, it should not persist", - ), - description="New description", - ), - ), - ], - ) - - instances = new_procedure.to_instances_write() - try: - created = workorder.unit_procedure.apply(new_procedure) - - assert len(created.nodes) == 2 - assert len(created.edges) == 1 - assert len(created.time_series) == 1 - finally: - cognite_client.data_modeling.instances.delete(instances.nodes.as_ids(), instances.edges.as_ids()) - cognite_client.time_series.delete(external_id=instances.time_series.as_external_ids(), ignore_unknown_ids=True) diff --git a/tests/test_unit/test_generator/conftest.py b/tests/test_unit/test_generator/conftest.py index 37ca88e7f..e567a84ed 100644 --- a/tests/test_unit/test_generator/conftest.py +++ b/tests/test_unit/test_generator/conftest.py @@ -1,32 +1,9 @@ import pytest from cognite.client import data_modeling as dm -from tests.constants import ( - EQUIPMENT_UNIT_SDK, - PUMP_SDK, -) +from tests.constants import PUMP_SDK @pytest.fixture(scope="session") def pump_model() -> dm.DataModel[dm.View]: return PUMP_SDK.load_data_model() - - -@pytest.fixture(scope="session") -def equipment_unit_model() -> dm.DataModel[dm.View]: - return EQUIPMENT_UNIT_SDK.load_data_model() - - -@pytest.fixture(scope="session") -def unit_procedure_view(equipment_unit_model: dm.DataModel) -> dm.View: - return next(v for v in equipment_unit_model.views if v.name == "UnitProcedure") - - -@pytest.fixture(scope="session") -def equipment_module_view(equipment_unit_model: dm.DataModel) -> dm.View: - return next(v for v in equipment_unit_model.views if v.name == "EquipmentModule") - - -@pytest.fixture(scope="session") -def start_end_time_view(equipment_unit_model: dm.DataModel) -> dm.View: - return next(v for v in equipment_unit_model.views if v.name == "StartEndTime") diff --git a/tests/test_unit/test_generator/test_sdk_generator_equipment_unit_sdk.py b/tests/test_unit/test_generator/test_sdk_generator_equipment_unit_sdk.py deleted file mode 100644 index a01df7015..000000000 --- a/tests/test_unit/test_generator/test_sdk_generator_equipment_unit_sdk.py +++ /dev/null @@ -1,191 +0,0 @@ -from __future__ import annotations - -import difflib - -import pytest -from cognite.client import data_modeling as dm - -from cognite.pygen._core.generators import APIGenerator, MultiAPIGenerator, SDKGenerator -from cognite.pygen._generator import CodeFormatter -from cognite.pygen.config import PygenConfig -from tests.constants import EQUIPMENT_UNIT_SDK, EquipmentSDKFiles - - -@pytest.fixture -def sdk_generator(equipment_unit_model: dm.DataModel[dm.View], pygen_config: PygenConfig) -> SDKGenerator: - return SDKGenerator( - EQUIPMENT_UNIT_SDK.top_level_package, - EQUIPMENT_UNIT_SDK.client_name, - equipment_unit_model, - config=pygen_config, - default_instance_space=EQUIPMENT_UNIT_SDK.instance_space, - ) - - -@pytest.fixture -def multi_api_generator(sdk_generator: SDKGenerator) -> MultiAPIGenerator: - return sdk_generator._multi_api_generator - - -@pytest.fixture -def unit_procedure_api_generator(multi_api_generator: MultiAPIGenerator, unit_procedure_view: dm.View) -> APIGenerator: - api_generator = multi_api_generator[unit_procedure_view.as_id()] - assert api_generator is not None, "Could not find API generator for unit procedure view" - return api_generator - - -@pytest.fixture -def equipment_module_api_generator( - multi_api_generator: MultiAPIGenerator, equipment_module_view: dm.View -) -> APIGenerator: - api_generator = multi_api_generator[equipment_module_view.as_id()] - assert api_generator is not None, "Could not find API generator for equipment module view" - return api_generator - - -@pytest.fixture -def start_end_time_api_generator(multi_api_generator: MultiAPIGenerator, start_end_time_view: dm.View) -> APIGenerator: - api_generator = multi_api_generator.api_by_type_by_view_id["edge"][start_end_time_view.as_id()] - assert api_generator is not None, "Could not find API generator for start end time view" - return api_generator - - -def test_generate_data_class_file_unit_procedure( - unit_procedure_api_generator: APIGenerator, pygen_config: PygenConfig, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.unit_procedure_data.read_text() - - # Act - actual = unit_procedure_api_generator.generate_data_class_file() - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_generate_data_class_start_end_time( - start_end_time_api_generator: APIGenerator, pygen_config: PygenConfig, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.start_end_time_data.read_text() - - # Act - actual = start_end_time_api_generator.generate_data_class_file() - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_generate_data_class_equipment_module( - equipment_module_api_generator: APIGenerator, pygen_config: PygenConfig, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.equipment_module_data.read_text() - - # Act - actual = equipment_module_api_generator.generate_data_class_file() - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_create_view_api_classes_equipment_module( - equipment_module_api_generator: APIGenerator, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.equipment_api.read_text() - - # Act - actual = equipment_module_api_generator.generate_api_file(EQUIPMENT_UNIT_SDK.client_name) - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_generate_equipment_module_sensor_value_api( - equipment_module_api_generator: APIGenerator, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.equipment_module_sensor_value_api.read_text() - - # Act - _, actual = next(equipment_module_api_generator.generate_timeseries_api_files(EQUIPMENT_UNIT_SDK.client_name)) - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_create_view_api_classes_unit_procedure( - unit_procedure_api_generator: APIGenerator, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.unit_procedure_api.read_text() - - # Act - actual = unit_procedure_api_generator.generate_api_file(EQUIPMENT_UNIT_SDK.client_name) - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_create_view_api_classes_unit_procedure_query( - unit_procedure_api_generator: APIGenerator, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.unit_procedure_query.read_text() - - # Act - actual = unit_procedure_api_generator.generate_api_query_file(EQUIPMENT_UNIT_SDK.client_name) - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_create_view_api_classes_unit_procedure_work_units( - unit_procedure_api_generator: APIGenerator, code_formatter: CodeFormatter -): - # Arrange - expected = EquipmentSDKFiles.unit_procedure_work_units.read_text() - - # Act - - actual_by_file_name = { - filename: content - for filename, content in unit_procedure_api_generator.generate_edge_api_files(EQUIPMENT_UNIT_SDK.client_name) - } - actual = actual_by_file_name[EquipmentSDKFiles.unit_procedure_work_units.stem] - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_generate_data_class_init_file(multi_api_generator: MultiAPIGenerator, code_formatter: CodeFormatter): - # Arrange - expected = EquipmentSDKFiles.data_init.read_text() - - # Act - actual = multi_api_generator.generate_data_classes_init_file() - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected - - -def test_create_api_client(sdk_generator: SDKGenerator, code_formatter: CodeFormatter): - # Arrange - expected = EquipmentSDKFiles.client.read_text() - - # Act - actual = sdk_generator._generate_api_client_file() - actual = code_formatter.format_code(actual) - - # Assert - assert actual == expected, "\n".join(difflib.unified_diff(expected.splitlines(), actual.splitlines()))