diff --git a/cognite/client/data_classes/_base.py b/cognite/client/data_classes/_base.py index 7f5b57be1..5fe525ea1 100644 --- a/cognite/client/data_classes/_base.py +++ b/cognite/client/data_classes/_base.py @@ -141,7 +141,7 @@ def dump_yaml(self) -> str: str: A YAML formatted string representing the instance. """ yaml = local_import("yaml") - return yaml.dump(self.dump(camel_case=True), sort_keys=False) + return yaml.safe_dump(self.dump(camel_case=True), sort_keys=False) @final @classmethod @@ -324,7 +324,7 @@ def dump_yaml(self) -> str: str: A YAML formatted string representing the instances. """ yaml = local_import("yaml") - return yaml.dump(self.dump(camel_case=True), sort_keys=False) + return yaml.safe_dump(self.dump(camel_case=True), sort_keys=False) def get(self, id: int | None = None, external_id: str | None = None) -> T_CogniteResource | None: """Get an item from this list by id or external_id. diff --git a/cognite/client/utils/_auxiliary.py b/cognite/client/utils/_auxiliary.py index 28a401234..1b6e91d76 100644 --- a/cognite/client/utils/_auxiliary.py +++ b/cognite/client/utils/_auxiliary.py @@ -16,6 +16,7 @@ from urllib.parse import quote from cognite.client.utils import _json +from cognite.client.utils._importing import local_import from cognite.client.utils._text import ( convert_all_keys_to_camel_case, convert_all_keys_to_snake_case, @@ -83,8 +84,7 @@ def fast_dict_load( def load_yaml_or_json(resource: str) -> Any: try: - import yaml - + yaml = local_import("yaml") return yaml.safe_load(resource) except ImportError: return _json.loads(resource) diff --git a/docs/source/conftest.py b/docs/source/conftest.py index 4bbdbada1..fabbae042 100644 --- a/docs/source/conftest.py +++ b/docs/source/conftest.py @@ -52,7 +52,7 @@ def client_data() -> dict[str, Any]: @pytest.fixture def quickstart_client_config_file(monkeypatch, client_data): def read_text(*args, **kwargs): - return yaml.dump(client_data) + return yaml.safe_dump(client_data) monkeypatch.setattr(Path, "read_text", read_text) diff --git a/tests/tests_unit/test_base.py b/tests/tests_unit/test_base.py index 03b411e19..f5c7ecf92 100644 --- a/tests/tests_unit/test_base.py +++ b/tests/tests_unit/test_base.py @@ -8,7 +8,6 @@ from unittest.mock import MagicMock import pytest -import yaml from cognite.client import ClientConfig, CogniteClient from cognite.client.credentials import Token @@ -306,8 +305,7 @@ def test_yaml_serialize(self, cognite_object_subclass: type[CogniteObject], cogn seed=65, cognite_client=cognite_mock_client_placeholder ).create_instance(cognite_object_subclass) - dumped = instance.dump(camel_case=True) - yaml_serialised = yaml.safe_dump(dumped) + yaml_serialised = instance.dump_yaml() loaded = instance.load(yaml_serialised, cognite_client=cognite_mock_client_placeholder) assert loaded.dump() == instance.dump() diff --git a/tests/tests_unit/test_data_classes/test_data_models/test_queries.py b/tests/tests_unit/test_data_classes/test_data_models/test_queries.py index 8d8bd6df6..bc87d257d 100644 --- a/tests/tests_unit/test_data_classes/test_data_models/test_queries.py +++ b/tests/tests_unit/test_data_classes/test_data_models/test_queries.py @@ -1,3 +1,4 @@ +import textwrap from collections.abc import Iterator from typing import Any @@ -7,6 +8,7 @@ from cognite.client.data_classes import filters as f from cognite.client.data_classes.data_modeling import InstanceSort, ViewId from cognite.client.data_classes.data_modeling import query as q +from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteFile def result_set_expression_load_and_dump_equals_data() -> Iterator[ParameterSet]: @@ -144,37 +146,38 @@ def test_dump(self, raw_data: dict, loaded: q.Select) -> None: def query_load_yaml_data() -> Iterator[ParameterSet]: - raw_yaml = """with: - airplanes: - nodes: - filter: - equals: - property: ["node", "externalId"] - value: {"parameter": "airplaneExternalId"} - chainTo: destination - direction: outwards - limit: 1 - lands_in_airports: - edges: - from: airplanes - maxDistance: 1 - direction: outwards - filter: - equals: - property: ["edge", "type"] - value: ["aviation", "lands-in"] - chainTo: destination - airports: - nodes: - chainTo: destination - direction: outwards - from: lands_in_airports -parameters: - airplaneExternalId: myFavouriteAirplane -select: - airplanes: {} - airports: {} -""" + raw_yaml = """\ + with: + airplanes: + nodes: + filter: + equals: + property: ["node", "externalId"] + value: {"parameter": "airplaneExternalId"} + chainTo: destination + direction: outwards + limit: 1 + lands_in_airports: + edges: + from: airplanes + maxDistance: 1 + direction: outwards + filter: + equals: + property: ["edge", "type"] + value: ["aviation", "lands-in"] + chainTo: destination + airports: + nodes: + chainTo: destination + direction: outwards + from: lands_in_airports + parameters: + airplaneExternalId: myFavouriteAirplane + select: + airplanes: {} + airports: {} + """ expected = q.Query( with_={ "airplanes": q.NodeResultSetExpression( @@ -191,34 +194,35 @@ def query_load_yaml_data() -> Iterator[ParameterSet]: parameters={"airplaneExternalId": "myFavouriteAirplane"}, select={"airplanes": q.Select(), "airports": q.Select()}, ) - yield pytest.param(raw_yaml, expected, id="Documentation Example") - - raw_yaml = """with: - movies: - nodes: - filter: - equals: - property: - - IntegrationTestsImmutable - - Movie/2 - - releaseYear - value: 1994 - chainTo: destination - direction: outwards -select: - movies: - sources: - - source: - space: IntegrationTestsImmutable - externalId: Movie - version: '2' - type: view - properties: - - title - - releaseYear -cursors: - movies: Z0FBQUFBQmtwc0RxQmducHpsWFd6VnZFdWwyWnFJbmxWS1BlT -""" + yield pytest.param(textwrap.dedent(raw_yaml), expected, id="Documentation Example") + + raw_yaml = """\ + with: + movies: + nodes: + filter: + equals: + property: + - IntegrationTestsImmutable + - Movie/2 + - releaseYear + value: 1994 + chainTo: destination + direction: outwards + select: + movies: + sources: + - source: + space: IntegrationTestsImmutable + externalId: Movie + version: '2' + type: view + properties: + - title + - releaseYear + cursors: + movies: Z0FBQUFBQmtwc0RxQmducHpsWFd6VnZFdWwyWnFJbmxWS1BlT + """ movie_id = ViewId(space="IntegrationTestsImmutable", external_id="Movie", version="2") movies_released_1994 = q.NodeResultSetExpression( filter=f.Equals(list(movie_id.as_property_ref("releaseYear")), 1994) @@ -228,7 +232,26 @@ def query_load_yaml_data() -> Iterator[ParameterSet]: select={"movies": q.Select([q.SourceSelector(movie_id, ["title", "releaseYear"])])}, cursors={"movies": "Z0FBQUFBQmtwc0RxQmducHpsWFd6VnZFdWwyWnFJbmxWS1BlT"}, ) - yield pytest.param(raw_yaml, expected, id="Example with cursors") + yield pytest.param(textwrap.dedent(raw_yaml), expected, id="Example with cursors") + + +@pytest.fixture +def query_with_non_yaml_native_data_classes() -> q.Query: + # YAML doesn't support tuple, so we want to test that yaml tags e.g. !!python/tuple does not end + # up in the output when dumping a Query object. + return q.Query( + with_={ + "files": q.NodeResultSetExpression( + filter=f.Exists(CogniteFile.get_source().as_property_ref("category")), + limit=None, + ), + "categories": q.NodeResultSetExpression( + from_="files", + through=CogniteFile.get_source().as_property_ref("category"), + ), + }, + select={"categories": q.Select()}, + ) class TestQuery: @@ -236,3 +259,17 @@ class TestQuery: def test_load_yaml(self, raw_data: str, expected: q.Query) -> None: actual = q.Query.load_yaml(raw_data) assert actual.dump(camel_case=True) == expected.dump(camel_case=True) + + def test_dump_yaml_no_tags(self, query_with_non_yaml_native_data_classes: q.Query) -> None: + query = query_with_non_yaml_native_data_classes + dumped = query.dump_yaml() + assert "!!python/tuple" not in dumped + + # Load will now load tuple as list: + loaded = q.Query.load_yaml(dumped) + assert loaded != query + + # ...re-dump-loading should be equal to the first loaded + dumped_again = loaded.dump_yaml() + assert "!!python/tuple" not in dumped_again + assert q.Query.load_yaml(dumped_again) == loaded diff --git a/tests/tests_unit/test_data_classes/test_hosted_extractors/test_jobs.py b/tests/tests_unit/test_data_classes/test_hosted_extractors/test_jobs.py index 864023bf3..e852ecba5 100644 --- a/tests/tests_unit/test_data_classes/test_hosted_extractors/test_jobs.py +++ b/tests/tests_unit/test_data_classes/test_hosted_extractors/test_jobs.py @@ -1,3 +1,5 @@ +import textwrap + import yaml from cognite.client.data_classes.hosted_extractors.jobs import Job @@ -5,19 +7,22 @@ class TestJob: def test_load_yaml_dump_unknown_config(self) -> None: - raw_yaml = """externalId: myJob -sourceId: my_eventhub -destinationId: EventHubTarget -targetStatus: running -status: running -createdTime: 123 -lastUpdatedTime: 1234 -format: - type: value - encoding: utf16 - compression: gzip -config: - some: new_config - that: has not been seen before""" - + raw_yaml = textwrap.dedent( + """\ + externalId: myJob + sourceId: my_eventhub + destinationId: EventHubTarget + targetStatus: running + status: running + createdTime: 123 + lastUpdatedTime: 1234 + format: + type: value + encoding: utf16 + compression: gzip + config: + some: new_config + that: has not been seen before + """ + ) assert Job.load(raw_yaml).dump() == yaml.safe_load(raw_yaml)