Skip to content

Commit

Permalink
remove yaml tags for tuple in dump_yaml methods (#1992)
Browse files Browse the repository at this point in the history
  • Loading branch information
haakonvt authored Oct 29, 2024
1 parent b971d77 commit ee2320e
Show file tree
Hide file tree
Showing 6 changed files with 123 additions and 83 deletions.
4 changes: 2 additions & 2 deletions cognite/client/data_classes/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions cognite/client/utils/_auxiliary.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion docs/source/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
4 changes: 1 addition & 3 deletions tests/tests_unit/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
157 changes: 97 additions & 60 deletions tests/tests_unit/test_data_classes/test_data_models/test_queries.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import textwrap
from collections.abc import Iterator
from typing import Any

Expand All @@ -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]:
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -228,11 +232,44 @@ 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:
@pytest.mark.parametrize("raw_data, expected", list(query_load_yaml_data()))
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
Original file line number Diff line number Diff line change
@@ -1,23 +1,28 @@
import textwrap

import yaml

from cognite.client.data_classes.hosted_extractors.jobs import Job


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)

0 comments on commit ee2320e

Please sign in to comment.