From 827961e4aa1bd6a28d2c408afd5ea571e3aa34c9 Mon Sep 17 00:00:00 2001 From: Austin G <19922895+austinlg96@users.noreply.github.com> Date: Tue, 25 Apr 2023 16:30:55 -0600 Subject: [PATCH] Prettier YAML output. --- nrel/hive/config/hive_config.py | 2 +- nrel/hive/custom_yaml/__init__.py | 1 + nrel/hive/custom_yaml/custom_yaml.py | 70 +++++++++ .../charging_search_type.py | 7 + nrel/hive/model/sim_time.py | 8 +- .../model/vehicle/schedules/schedule_type.py | 7 + nrel/hive/reporting/report_type.py | 16 +++ tests/test_charging_search_type.py | 13 ++ tests/test_custom_yaml.py | 136 ++++++++++++++++++ tests/test_report_type.py | 32 +++++ tests/test_schedule_type.py | 10 ++ tests/test_sim_time.py | 10 ++ 12 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 nrel/hive/custom_yaml/__init__.py create mode 100644 nrel/hive/custom_yaml/custom_yaml.py create mode 100644 tests/test_charging_search_type.py create mode 100644 tests/test_custom_yaml.py create mode 100644 tests/test_report_type.py create mode 100644 tests/test_schedule_type.py create mode 100644 tests/test_sim_time.py diff --git a/nrel/hive/config/hive_config.py b/nrel/hive/config/hive_config.py index 51852d76..a6ed9702 100644 --- a/nrel/hive/config/hive_config.py +++ b/nrel/hive/config/hive_config.py @@ -8,7 +8,6 @@ from typing import NamedTuple, Dict, Union, Tuple, Optional import pkg_resources -import yaml from nrel.hive.config.config_builder import ConfigBuilder from nrel.hive.config.dispatcher_config import DispatcherConfig @@ -16,6 +15,7 @@ from nrel.hive.config.input import Input from nrel.hive.config.network import Network from nrel.hive.config.sim import Sim +from nrel.hive.custom_yaml import custom_yaml as yaml from nrel.hive.util import fs log = logging.getLogger(__name__) diff --git a/nrel/hive/custom_yaml/__init__.py b/nrel/hive/custom_yaml/__init__.py new file mode 100644 index 00000000..35f9ca85 --- /dev/null +++ b/nrel/hive/custom_yaml/__init__.py @@ -0,0 +1 @@ +from nrel.hive.custom_yaml.custom_yaml import custom_yaml diff --git a/nrel/hive/custom_yaml/custom_yaml.py b/nrel/hive/custom_yaml/custom_yaml.py new file mode 100644 index 00000000..6a8f3349 --- /dev/null +++ b/nrel/hive/custom_yaml/custom_yaml.py @@ -0,0 +1,70 @@ +import logging +from pathlib import PurePath +from typing import Any, Union + +import yaml + +from nrel.hive.dispatcher.instruction_generator.charging_search_type import ChargingSearchType +from nrel.hive.model.sim_time import SimTime +from nrel.hive.model.vehicle.schedules.schedule_type import ScheduleType +from nrel.hive.reporting.report_type import ReportType + +log = logging.getLogger(__name__) + +custom_yaml = yaml + +# This tag is not written to the file during serialization because interpretation is implicit during YAML deserialization. +YAML_STR_TAG = "tag:yaml.org,2002:str" + + +# Handling stdlib objects that should be represented by list(obj). +# Prefer to handle classes within their own definition and then register below. +def convert_to_unsorted_list(dumper: custom_yaml.Dumper, obj: tuple): + """Patches PyYAML representation for an object so that it is treated as a YAML list. Avoids an explicit YAML tag.""" + return dumper.represent_list(list(obj)) + + +custom_yaml.add_representer(data_type=tuple, representer=convert_to_unsorted_list) + + +# Handling stdlib objects that should be represented by sorted(list(obj)). +# Prefer to handle classes within their own definition and then register below. +def convert_to_sorted_list(dumper: custom_yaml.Dumper, obj: set): + """Patches PyYAML representation for an object so that it is treated as a YAML list. Avoids an explicit YAML tag.""" + return dumper.represent_list(sorted(list(obj))) + + +custom_yaml.add_representer(data_type=set, representer=convert_to_sorted_list) + + +# Handling stdlib objects that should be represented as str(obj). +# Prefer to handle classes within their own definition and then register below. +def convert_to_str(dumper: custom_yaml.Dumper, path: PurePath): + """Patches PyYAML representation for an object so that it is treated as a YAML str. Avoids an explicit YAML tag.""" + return dumper.represent_scalar(tag=YAML_STR_TAG, value=str(path)) + + +custom_yaml.add_multi_representer(data_type=PurePath, multi_representer=convert_to_str) + + +# Registering explicit/specific representers that are kept withing their classes. +custom_yaml.add_representer( + data_type=ChargingSearchType, representer=ChargingSearchType.yaml_representer +) +custom_yaml.add_representer(data_type=ReportType, representer=ReportType.yaml_representer) +custom_yaml.add_representer(data_type=ScheduleType, representer=ScheduleType.yaml_representer) +custom_yaml.add_representer(data_type=SimTime, representer=SimTime.yaml_representer) + + +# Fallback to str() representation for any child of `object`. +# Raise a warning to alert the user that implicit serialization was done. +# Does not appear to work for built-in types that PyYaml has special serializers for. +def generic_representer(dumper: custom_yaml.Dumper, obj: object): + """Serializes arbitrary objects to strs.""" + log.warning(f"{obj.__class__} object was implicity serialized with `str(obj)`.") + tag = YAML_STR_TAG + val = str(obj) + return dumper.represent_scalar(tag=tag, value=val) + + +custom_yaml.add_multi_representer(data_type=object, multi_representer=generic_representer) diff --git a/nrel/hive/dispatcher/instruction_generator/charging_search_type.py b/nrel/hive/dispatcher/instruction_generator/charging_search_type.py index 4be6b43c..4fdb3576 100644 --- a/nrel/hive/dispatcher/instruction_generator/charging_search_type.py +++ b/nrel/hive/dispatcher/instruction_generator/charging_search_type.py @@ -1,7 +1,10 @@ from __future__ import annotations from enum import Enum +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import yaml class ChargingSearchType(Enum): NEAREST_SHORTEST_QUEUE = 1 @@ -26,3 +29,7 @@ def from_string(string: str) -> ChargingSearchType: raise NameError( f"charging search type {string} is not known, must be one of {valid_names}" ) + + @staticmethod + def yaml_representer(dumper: yaml.Dumper, o: "ChargingSearchType"): + return dumper.represent_scalar(tag = 'tag:yaml.org,2002:str', value = o.name.lower()) \ No newline at end of file diff --git a/nrel/hive/model/sim_time.py b/nrel/hive/model/sim_time.py index d1defa12..95bcccb8 100644 --- a/nrel/hive/model/sim_time.py +++ b/nrel/hive/model/sim_time.py @@ -1,10 +1,12 @@ from __future__ import annotations from datetime import datetime, time -from typing import Union +from typing import TYPE_CHECKING, Union from nrel.hive.util.exception import TimeParseError +if TYPE_CHECKING: + import yaml class SimTime(int): ERROR_MSG = ( @@ -65,3 +67,7 @@ def as_epoch_time(self) -> int: def as_iso_time(self) -> str: return datetime.utcfromtimestamp(int(self)).isoformat() + + @staticmethod + def yaml_representer(dumper: yaml.Dumper, o: "SimTime"): + return dumper.represent_scalar(tag = 'tag:yaml.org,2002:str', value = o.as_iso_time()) \ No newline at end of file diff --git a/nrel/hive/model/vehicle/schedules/schedule_type.py b/nrel/hive/model/vehicle/schedules/schedule_type.py index 325481ab..2ccafb0c 100644 --- a/nrel/hive/model/vehicle/schedules/schedule_type.py +++ b/nrel/hive/model/vehicle/schedules/schedule_type.py @@ -1,7 +1,10 @@ from __future__ import annotations from enum import Enum +from typing import TYPE_CHECKING +if TYPE_CHECKING: + import yaml class ScheduleType(Enum): TIME_RANGE = 0 @@ -24,3 +27,7 @@ def from_string(string: str) -> ScheduleType: raise NameError( f"schedule type '{string}' is not known, must be one of {{{valid_names}}}" ) + + @staticmethod + def yaml_representer(dumper: yaml.Dumper, o: "ScheduleType"): + return dumper.represent_scalar(tag = 'tag:yaml.org,2002:str', value = o.name.lower()) \ No newline at end of file diff --git a/nrel/hive/reporting/report_type.py b/nrel/hive/reporting/report_type.py index cc3e244e..236f152e 100644 --- a/nrel/hive/reporting/report_type.py +++ b/nrel/hive/reporting/report_type.py @@ -1,6 +1,10 @@ from __future__ import annotations from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import yaml class ReportType(Enum): @@ -43,3 +47,15 @@ def from_string(cls, s: str) -> ReportType: return values[s] except KeyError: raise KeyError(f"{s} not a valid report type.") + + @staticmethod + def yaml_representer(dumper: yaml.Dumper, o: "ReportType"): + return dumper.represent_scalar(tag="tag:yaml.org,2002:str", value=o.name.lower()) + + def __lt__(self, other: "ReportType"): + """Allows sorting an iterable of ReportType, in particular for deterministic serialization of `set[ReportType]`""" + if not isinstance(other, ReportType): + raise TypeError( + f"'<' not supported between instances of {type(self)} and {type(other)}" + ) + return self.name < other.name diff --git a/tests/test_charging_search_type.py b/tests/test_charging_search_type.py new file mode 100644 index 00000000..042c6f9a --- /dev/null +++ b/tests/test_charging_search_type.py @@ -0,0 +1,13 @@ +from unittest import TestCase + +from nrel.hive.custom_yaml import custom_yaml as yaml +from nrel.hive.dispatcher.instruction_generator.charging_search_type import ChargingSearchType + + +class TestChargingSearchType(TestCase): + def test_yaml_repr(self): + a = ChargingSearchType(1) + yaml.add_representer( + data_type=ChargingSearchType, representer=ChargingSearchType.yaml_representer + ) + self.assertEqual(yaml.dump(a), "nearest_shortest_queue\n...\n") diff --git a/tests/test_custom_yaml.py b/tests/test_custom_yaml.py new file mode 100644 index 00000000..0f22d112 --- /dev/null +++ b/tests/test_custom_yaml.py @@ -0,0 +1,136 @@ +import logging +from pathlib import Path, PurePath +from unittest import TestCase + +from nrel.hive.custom_yaml import custom_yaml as yaml +from nrel.hive.dispatcher.instruction_generator.charging_search_type import ChargingSearchType +from nrel.hive.model.sim_time import SimTime +from nrel.hive.model.vehicle.schedules.schedule_type import ScheduleType +from nrel.hive.reporting.report_type import ReportType + +BASIC_TUPLE = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9) +BASIC_LIST = list(BASIC_TUPLE) + +BASIC_SET = set([3, 2, 1, 0, 4, 6, 7, 5, 9, 8]) + + +class TestClass(object): + """An accessory class for use in tests without any YAML serialization methods.""" + + def __str__(self): + return "123abc123" + + +log = logging.getLogger(__name__) + + +class TestCustomYaml_Tuple(TestCase): + def test_not_tagged_YAML(self): + a = (1,) + self.assertNotEqual(yaml.dump(a)[0], "!", "YAML ! tag detected during serialization.") + + def test_tuple_like_list(self): + self.assertEqual( + yaml.dump(BASIC_LIST), + yaml.dump(BASIC_TUPLE), + "The custom YAML serializer is not treating a tuple like a YAML list.", + ) + + +class TestCustomYAML_Set(TestCase): + def test_not_tagged_YAML(self): + a = set([1]) + self.assertNotEqual(yaml.dump(a)[0], "!", "YAML ! tag detected during serialization.") + + def test_set_like_list(self): + self.assertEqual( + yaml.dump(BASIC_LIST), + yaml.dump(BASIC_SET), + "The custom YAML serializer is not treating a set like a YAML list.", + ) + + +class TestCustomYAML_PathLib(TestCase): + def test_purepath(self): + ppath = PurePath("./test/") + str_path = str(ppath) + self.assertEqual(yaml.dump(str_path), yaml.dump(ppath)) + + def test_path(self): + path = Path("./test/") + str_path = str(path) + self.assertEqual(yaml.dump(str_path), yaml.dump(path)) + + +class TestCustomYAML_ChargingSearchType(TestCase): + def test_not_tagged_YAML(self): + a = ChargingSearchType(1) + self.assertNotEqual(yaml.dump(a)[0], "!", "YAML ! tag detected during serialization.") + + def test_like_str(self): + a = ChargingSearchType(1) + self.assertEqual(yaml.dump(a), "nearest_shortest_queue\n...\n") + + +class TestCustomYAML_SimTime(TestCase): + def test_not_tagged_YAML(self): + a = SimTime.build(0) + self.assertNotEqual(yaml.dump(a)[0], "!", "YAML ! tag detected during serialization.") + + def test_like_iso_time_dynamic(self): + a = SimTime.build(0) + self.assertEqual(yaml.dump(a.as_iso_time()), yaml.dump(a)) + + def test_like_iso_time_static(self): + a = SimTime.build(0) + self.assertEqual(yaml.dump(a.as_iso_time()), "'1970-01-01T00:00:00'\n") + + +class TestCustomYAML_ScheduleType(TestCase): + def test_not_tagged_YAML(self): + a = ScheduleType(0) + self.assertNotEqual(yaml.dump(a)[0], "!", "YAML ! tag detected during serialization.") + + def test_like_str(self): + a = ScheduleType(0) + self.assertEqual(yaml.dump(a), "time_range\n...\n") + + +class TestCustomYAML_ReportType(TestCase): + def test_not_tagged_YAML(self): + a = ReportType(1) + self.assertNotEqual(yaml.dump(a)[0], "!", "YAML ! tag detected during serialization.") + + def test_like_str(self): + a = ReportType(1) + self.assertEqual(yaml.dump(a), "station_state\n...\n") + + +class TestCustomYAML_Generic(TestCase): + def test_not_tagged_YAML(self): + instance = TestClass() + self.assertNotEqual( + yaml.dump(instance)[0], "!", "YAML ! tag detected during serialization." + ) + + def test_warn_on_generic_serialization(self): + for c in (TestClass, lambda: range(10)): + instance = c() + + with self.assertLogs(level=logging.WARNING) as log_cm: + a = yaml.dump(instance) + print(a) + + warninglog = ( + "WARNING:nrel.hive.custom_yaml.custom_yaml:" + + f"{instance.__class__} object was implicity serialized with `str(obj)`." + ) + self.assertIn( + warninglog, + log_cm.output, + msg="WARNING entry in log was not present when implicitly serializing an object.", + ) + + def test_generic_serialization_like_str(self): + instance = TestClass() + self.assertEqual(yaml.dump(str(instance)), yaml.dump(instance)) diff --git a/tests/test_report_type.py b/tests/test_report_type.py new file mode 100644 index 00000000..cadb98ca --- /dev/null +++ b/tests/test_report_type.py @@ -0,0 +1,32 @@ +from unittest import TestCase + +from nrel.hive.custom_yaml import custom_yaml as yaml +from nrel.hive.reporting.report_type import ReportType + + +class TestReportType(TestCase): + def test_yaml_repr(self): + yaml.add_representer(data_type=ReportType, representer=ReportType.yaml_representer) + self.assertEqual("station_state\n...\n", yaml.dump(ReportType.from_string("station_state"))) + + def test_ReportType_ordering_dynamic(self): + members = [m for m in ReportType] + name_then_sort = sorted([m.name for m in members]) + sort_then_name = [m.name for m in sorted(members)] + self.assertEqual(name_then_sort, sort_then_name, "ReportType sorting invalid.") + + def test_ReportType_ordering_static(self): + a = ReportType.from_string("station_state") + b = ReportType.from_string("driver_state") + self.assertLess(b, a, "ReportType sorting invalid.") + + def test_ReportType_lt_raise(self): + class GenericClass: + @property + def name(self): + return "abc" + + a = ReportType.from_string("station_state") + b = GenericClass() + with self.assertRaises(TypeError): + x = a < b diff --git a/tests/test_schedule_type.py b/tests/test_schedule_type.py new file mode 100644 index 00000000..c1ba087b --- /dev/null +++ b/tests/test_schedule_type.py @@ -0,0 +1,10 @@ +from unittest import TestCase + +from nrel.hive.custom_yaml import custom_yaml as yaml +from nrel.hive.model.vehicle.schedules.schedule_type import ScheduleType + + +class TestSimTime(TestCase): + def test_yaml_repr(self): + yaml.add_representer(data_type=ScheduleType, representer=ScheduleType.yaml_representer) + self.assertEqual("time_range\n...\n", yaml.dump(ScheduleType(0))) diff --git a/tests/test_sim_time.py b/tests/test_sim_time.py new file mode 100644 index 00000000..a5c99bdd --- /dev/null +++ b/tests/test_sim_time.py @@ -0,0 +1,10 @@ +from unittest import TestCase + +from nrel.hive.custom_yaml import custom_yaml as yaml +from nrel.hive.model.sim_time import SimTime + + +class TestSimTime(TestCase): + def test_yaml_repr(self): + yaml.add_representer(data_type=SimTime, representer=SimTime.yaml_representer) + self.assertEqual("'1970-01-01T00:00:00'\n", yaml.dump(SimTime.build(0)))