From d3c954528597994078513e4407eb6a2329adf8a8 Mon Sep 17 00:00:00 2001 From: Spencer Brown Date: Thu, 23 Jan 2025 18:17:18 +1000 Subject: [PATCH] Add Element.get_wrap() --- docs/source/changelog.rst | 1 + src/srctools/dmx.py | 15 +++++++++++++++ tests/test_dmx.py | 36 ++++++++++++++++++++++++++---------- 3 files changed, 42 insertions(+), 10 deletions(-) diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst index 10aa20b4..7fd599ea 100644 --- a/docs/source/changelog.rst +++ b/docs/source/changelog.rst @@ -21,6 +21,7 @@ Version (dev) * Handle entities being added/removed during iteration of :py:meth:`VMF.search() `. * Share common strings in the engine database to save some space. * Fix saving the `PHYSCOLLIDE` BSP lump. +* Add :py:meth:`srctools.dmx.Element.get_wrap()`, allowing handling defaults more conveniently. * Make :py:attr:`EntityDef.kv `, :py:attr:`.inp ` and :py:attr:`.out ` views settable, improve behaviour. diff --git a/src/srctools/dmx.py b/src/srctools/dmx.py index b7baab2b..73869ad0 100644 --- a/src/srctools/dmx.py +++ b/src/srctools/dmx.py @@ -2051,6 +2051,21 @@ def clear(self) -> None: """Remove all attributes from the element.""" self._members.clear() + def get_wrap(self, name: str, default: Union[Attribute, ConvValue, Sequence[ConvValue]] = '', /) -> Attribute: + """Retrieve the specified attribute, or return a default. + + If the defalt is an Attribute it is returned unchanged, otherwise it is wrapped in + a temporary attribute. + """ + try: + return self._members[name.casefold()] + except KeyError: + if isinstance(default, Attribute): + return default + else: + typ, val = deduce_type(default) + return Attribute(name, typ, val) + def pop(self, name: str, default: Union[Attribute, ConvValue, Sequence[ConvValue]] = _UNSET) -> Attribute: """Remove the specified attribute and return it. diff --git a/tests/test_dmx.py b/tests/test_dmx.py index cd49ff85..ab574b53 100644 --- a/tests/test_dmx.py +++ b/tests/test_dmx.py @@ -1,5 +1,5 @@ """Test the datamodel exchange implementation.""" -from typing import cast +from typing import Any, List, Tuple, cast from collections.abc import Callable from io import BytesIO from pathlib import Path @@ -489,7 +489,7 @@ def test_binary_text_conversion(typ: ValueType, attr: str, value, binary: bytes, 'int', 'float', 'bool', 'str', 'bytes', 'time', 'color', 'vec2', 'vec3', 'vec4', 'angle', 'quaternion', 'matrix', ]) -def test_attr_array_constructor(typ: ValueType, iterable: list, expected: list) -> None: +def test_attr_array_constructor(typ: ValueType, iterable: List[Any], expected: List[Any]) -> None: """Test constructing an array attribute.""" expected_type = type(expected[0]) attr = Attribute.array('some_array', typ, iterable) @@ -515,7 +515,7 @@ def test_attr_array_elem_constructor() -> None: assert attr[2].val_elem is elem1 -deduce_type_tests = [ +deduce_type_tests: List[Tuple[Any, ValueType, Any]] = [ (5, ValueType.INT, 5), (5.0, ValueType.FLOAT, 5.0), (False, ValueType.BOOL, False), @@ -537,16 +537,16 @@ def test_attr_array_elem_constructor() -> None: ] -@pytest.mark.parametrize('input, val_type, output', deduce_type_tests) -def test_deduce_type_basic(input, val_type, output) -> None: +@pytest.mark.parametrize('input_val, val_type, output', deduce_type_tests) +def test_deduce_type_basic(input_val: Any, val_type: ValueType, output: Any) -> None: """Test type deduction behaves correctly.""" - [test_type, test_val] = deduce_type(input) + [test_type, test_val] = deduce_type(input_val) assert test_type is val_type assert type(test_val) is type(output) assert test_val == output -@pytest.mark.parametrize('input, val_type, output', [ +@pytest.mark.parametrize('input_val, val_type, output', [ # Add the above tests here too. ([inp, inp, inp], val_typ, [out, out, out]) for inp, val_typ, out in @@ -574,11 +574,11 @@ def test_deduce_type_basic(input, val_type, output) -> None: (collections.deque([1.0, 2.0, 3.0]), ValueType.FLOAT, [1.0, 2.0, 3.0]), (range(5), ValueType.INT, [0, 1, 2, 3, 4]), ]) -def test_deduce_type_array(input, val_type, output) -> None: +def test_deduce_type_array(input_val: List[Any], val_type: ValueType, output: List[Any]) -> None: """Test array deduction, and some special cases.""" - [test_type, test_arr] = deduce_type(input) + [test_type, test_arr] = deduce_type(input_val) assert test_type is val_type - assert len(input) == len(test_arr), repr(test_arr) + assert len(input_val) == len(test_arr), repr(test_arr) for i, (test, out) in enumerate(zip(test_arr, output)): assert type(test) is type(out), f'{i}: {test!r} != {out!r}' assert test == out, f'{i}: {test!r} != {out!r}' @@ -601,6 +601,22 @@ def test_deduce_type_adv() -> None: print(deduce_type(range(0))) +def test_elem_get() -> None: + """Test the behaviours that get various values.""" + elem = Element('Named', 'Type') + elem['keyName'] = attr = Attribute.string('keyName', 'hello') + assert elem['keyName'] is attr + with pytest.raises(KeyError): + elem['missing'] + assert elem.get_wrap('keyName') is attr + assert elem.get_wrap('keyName', 'default') is attr + assert elem.get_wrap('miSSing', 45) == Attribute.int('miSSing', 45) + assert elem.get_wrap('miSSing', True) == Attribute.bool('miSSing', True) + assert elem.get_wrap('miSSing', [1.0, 2.0]) == Attribute.array( + 'miSSing', ValueType.FLOAT, [1.0, 2.0], + ) + + def test_special_attr_name() -> None: """Test the special behaviour of the "name" attribute.