[PATCH v2 2/2] dts: clean up config types

[Luca Vizzarro <luca.vizzarro@arm.com>]

Clean up types used with the configuration classes, and use Self from
the newly added typing_extensions module.

Methods that instantiate their own class should be @classmethod instead
of @staticmethod.

Bugzilla ID: 1433

Signed-off-by: Luca Vizzarro <luca.vizzarro@arm.com>
Reviewed-by: Paul Szczepanek <paul.szczepanek@arm.com>
---
 dts/framework/config/__init__.py              | 47 ++++++++++---------
 .../traffic_generator/__init__.py             | 10 ++--
 dts/poetry.lock                               |  2 +-
 dts/pyproject.toml                            |  1 +
 4 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/dts/framework/config/__init__.py b/dts/framework/config/__init__.py
index 6b2ad2b16d..41b2fbc94d 100644
--- a/dts/framework/config/__init__.py
+++ b/dts/framework/config/__init__.py
@@ -42,6 +42,7 @@
 
 import warlock  # type: ignore[import-untyped]
 import yaml
+from typing_extensions import Self
 
 from framework.config.types import (
     BuildTargetConfigDict,
@@ -156,8 +157,8 @@ class PortConfig:
     peer_node: str
     peer_pci: str
 
-    @staticmethod
-    def from_dict(node: str, d: PortConfigDict) -> "PortConfig":
+    @classmethod
+    def from_dict(cls, node: str, d: PortConfigDict) -> Self:
         """A convenience method that creates the object from fewer inputs.
 
         Args:
@@ -167,7 +168,7 @@ def from_dict(node: str, d: PortConfigDict) -> "PortConfig":
         Returns:
             The port configuration instance.
         """
-        return PortConfig(node=node, **d)
+        return cls(node=node, **d)
 
 
 @dataclass(slots=True, frozen=True)
@@ -183,7 +184,7 @@ class TrafficGeneratorConfig:
     traffic_generator_type: TrafficGeneratorType
 
     @staticmethod
-    def from_dict(d: TrafficGeneratorConfigDict) -> "ScapyTrafficGeneratorConfig":
+    def from_dict(d: TrafficGeneratorConfigDict) -> "TrafficGeneratorConfig":
         """A convenience method that produces traffic generator config of the proper type.
 
         Args:
@@ -312,7 +313,7 @@ class TGNodeConfiguration(NodeConfiguration):
         traffic_generator: The configuration of the traffic generator present on the TG node.
     """
 
-    traffic_generator: ScapyTrafficGeneratorConfig
+    traffic_generator: TrafficGeneratorConfig
 
 
 @dataclass(slots=True, frozen=True)
@@ -356,8 +357,8 @@ class BuildTargetConfiguration:
     compiler_wrapper: str
     name: str
 
-    @staticmethod
-    def from_dict(d: BuildTargetConfigDict) -> "BuildTargetConfiguration":
+    @classmethod
+    def from_dict(cls, d: BuildTargetConfigDict) -> Self:
         r"""A convenience method that processes the inputs before creating an instance.
 
         `arch`, `os`, `cpu` and `compiler` are converted to :class:`Enum`\s and
@@ -369,7 +370,7 @@ def from_dict(d: BuildTargetConfigDict) -> "BuildTargetConfiguration":
         Returns:
             The build target configuration instance.
         """
-        return BuildTargetConfiguration(
+        return cls(
             arch=Architecture(d["arch"]),
             os=OS(d["os"]),
             cpu=CPUType(d["cpu"]),
@@ -407,10 +408,11 @@ class TestSuiteConfig:
     test_suite: str
     test_cases: list[str]
 
-    @staticmethod
+    @classmethod
     def from_dict(
+        cls,
         entry: str | TestSuiteConfigDict,
-    ) -> "TestSuiteConfig":
+    ) -> Self:
         """Create an instance from two different types.
 
         Args:
@@ -420,9 +422,9 @@ def from_dict(
             The test suite configuration instance.
         """
         if isinstance(entry, str):
-            return TestSuiteConfig(test_suite=entry, test_cases=[])
+            return cls(test_suite=entry, test_cases=[])
         elif isinstance(entry, dict):
-            return TestSuiteConfig(test_suite=entry["suite"], test_cases=entry["cases"])
+            return cls(test_suite=entry["suite"], test_cases=entry["cases"])
         else:
             raise TypeError(f"{type(entry)} is not valid for a test suite config.")
 
@@ -454,11 +456,12 @@ class ExecutionConfiguration:
     traffic_generator_node: TGNodeConfiguration
     vdevs: list[str]
 
-    @staticmethod
+    @classmethod
     def from_dict(
+        cls,
         d: ExecutionConfigDict,
-        node_map: dict[str, Union[SutNodeConfiguration | TGNodeConfiguration]],
-    ) -> "ExecutionConfiguration":
+        node_map: dict[str, SutNodeConfiguration | TGNodeConfiguration],
+    ) -> Self:
         """A convenience method that processes the inputs before creating an instance.
 
         The build target and the test suite config are transformed into their respective objects.
@@ -494,7 +497,7 @@ def from_dict(
         vdevs = (
             d["system_under_test_node"]["vdevs"] if "vdevs" in d["system_under_test_node"] else []
         )
-        return ExecutionConfiguration(
+        return cls(
             build_targets=build_targets,
             perf=d["perf"],
             func=d["func"],
@@ -505,7 +508,7 @@ def from_dict(
             vdevs=vdevs,
         )
 
-    def copy_and_modify(self, **kwargs) -> "ExecutionConfiguration":
+    def copy_and_modify(self, **kwargs) -> Self:
         """Create a shallow copy with any of the fields modified.
 
         The only new data are those passed to this method.
@@ -525,7 +528,7 @@ def copy_and_modify(self, **kwargs) -> "ExecutionConfiguration":
             else:
                 new_config[field.name] = getattr(self, field.name)
 
-        return ExecutionConfiguration(**new_config)
+        return type(self)(**new_config)
 
 
 @dataclass(slots=True, frozen=True)
@@ -541,8 +544,8 @@ class Configuration:
 
     executions: list[ExecutionConfiguration]
 
-    @staticmethod
-    def from_dict(d: ConfigurationDict) -> "Configuration":
+    @classmethod
+    def from_dict(cls, d: ConfigurationDict) -> Self:
         """A convenience method that processes the inputs before creating an instance.
 
         Build target and test suite config are transformed into their respective objects.
@@ -555,7 +558,7 @@ def from_dict(d: ConfigurationDict) -> "Configuration":
         Returns:
             The whole configuration instance.
         """
-        nodes: list[Union[SutNodeConfiguration | TGNodeConfiguration]] = list(
+        nodes: list[SutNodeConfiguration | TGNodeConfiguration] = list(
             map(NodeConfiguration.from_dict, d["nodes"])
         )
         assert len(nodes) > 0, "There must be a node to test"
@@ -567,7 +570,7 @@ def from_dict(d: ConfigurationDict) -> "Configuration":
             map(ExecutionConfiguration.from_dict, d["executions"], [node_map for _ in d])
         )
 
-        return Configuration(executions=executions)
+        return cls(executions=executions)
 
 
 def load_config(config_file_path: Path) -> Configuration:
diff --git a/dts/framework/testbed_model/traffic_generator/__init__.py b/dts/framework/testbed_model/traffic_generator/__init__.py
index 0eaf0355cd..03e57a77fc 100644
--- a/dts/framework/testbed_model/traffic_generator/__init__.py
+++ b/dts/framework/testbed_model/traffic_generator/__init__.py
@@ -16,7 +16,7 @@
 
 # pylama:ignore=W0611
 
-from framework.config import ScapyTrafficGeneratorConfig, TrafficGeneratorType
+from framework.config import ScapyTrafficGeneratorConfig, TrafficGeneratorConfig
 from framework.exception import ConfigurationError
 from framework.testbed_model.node import Node
 
@@ -28,7 +28,7 @@
 
 
 def create_traffic_generator(
-    tg_node: Node, traffic_generator_config: ScapyTrafficGeneratorConfig
+    tg_node: Node, traffic_generator_config: TrafficGeneratorConfig
 ) -> CapturingTrafficGenerator:
     """The factory function for creating traffic generator objects from the test run configuration.
 
@@ -39,10 +39,10 @@ def create_traffic_generator(
     Returns:
         A traffic generator capable of capturing received packets.
     """
-    match traffic_generator_config.traffic_generator_type:
-        case TrafficGeneratorType.SCAPY:
+    match traffic_generator_config:
+        case ScapyTrafficGeneratorConfig():
             return ScapyTrafficGenerator(tg_node, traffic_generator_config)
         case _:
             raise ConfigurationError(
-                "Unknown traffic generator: {traffic_generator_config.traffic_generator_type}"
+                f"Unknown traffic generator: {traffic_generator_config.traffic_generator_type}"
             )
diff --git a/dts/poetry.lock b/dts/poetry.lock
index df9cecb7e0..5f8fa03933 100644
--- a/dts/poetry.lock
+++ b/dts/poetry.lock
@@ -853,4 +853,4 @@ jsonschema = ">=4,<5"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "1572cb14f0bf88cddc6b6b225312ad468d01916b32067d17a7776af76c6d466c"
+content-hash = "4af4dd49c59e5bd6ed99e8c19c6756aaf00125339d26cfad2ef98551dc765f8b"
diff --git a/dts/pyproject.toml b/dts/pyproject.toml
index 05c91ef9be..a160d2fa02 100644
--- a/dts/pyproject.toml
+++ b/dts/pyproject.toml
@@ -26,6 +26,7 @@ types-PyYAML = "^6.0.8"
 fabric = "^2.7.1"
 scapy = "^2.5.0"
 pydocstyle = "6.1.1"
+typing-extensions = "^4.11.0"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.10.0"
-- 
2.34.1