Merge pull request #1 from xenanetworks/dev

Switch to pydantic v2

Author: leonardhyu

plugin2544/dataset.py

@@ -1,5 +1,5 @@
-from typing import Any, List, Tuple, Dict
-from pydantic import BaseModel, validator
+from typing import Any, List, Tuple, Dict, Annotated
+from pydantic import BaseModel, field_validator, ValidationInfo, Field
 from .utils import exceptions, constants as const
 from .model.m_test_config import TestConfigModel
 from .model.m_test_type_config import TestTypesConfiguration
@@ -11,9 +11,9 @@
 
 
 class PluginModel2544(BaseModel):  # Main Model
-    test_configuration: TestConfigModel
+    test_configuration: Annotated[TestConfigModel, Field(validate_default=True)]
     protocol_segments: List[ProtocolSegmentProfileConfig]
-    ports_configuration: PortConfType
+    ports_configuration: Annotated[PortConfType, Field(validate_default=True)]
     test_types_configuration: TestTypesConfiguration
 
     def set_ports_rx_tx_type(self) -> None:
@@ -45,33 +45,33 @@ def __init__(self, **data: Dict[str, Any]) -> None:
         self.check_port_groups_and_peers()
         self.set_profile()
 
-    @validator("ports_configuration", always=True)
-    def check_ip_properties(cls, v: "PortConfType", values) -> "PortConfType":
-        pro_map = {v.id: v.protocol_version for v in values['protocol_segments']}
-        for i, port_config in enumerate(v):
+    
+    @field_validator("ports_configuration")
+    def check_ip_properties(cls, value: "PortConfType", info: ValidationInfo) -> "PortConfType":
+        pro_map = {v.id: v.protocol_version for v in info.data['protocol_segments']}
+        for i, port_config in enumerate(value):
             if port_config.protocol_segment_profile_id not in pro_map:
                 raise exceptions.PSPMissing()
             if (
                 pro_map[port_config.protocol_segment_profile_id].is_l3
                 and (not port_config.ip_address or port_config.ip_address.address.is_empty)
             ):
                 raise exceptions.IPAddressMissing()
-        return v
+        return value
 
-    @validator("ports_configuration", always=True)
-    def check_port_count(
-        cls, v: "PortConfType", values: Dict[str, Any]
-    ) -> "PortConfType":
+    
+    @field_validator("ports_configuration")
+    def check_port_count(cls, value: "PortConfType", info: ValidationInfo) -> "PortConfType":
         require_ports = 2
-        if "test_configuration" in values:
-            topology: const.TestTopology = values[
+        if "test_configuration" in info.data:
+            topology: const.TestTopology = info.data[
                 "test_configuration"
             ].topology_config.topology
             if topology.is_pair_topology:
                 require_ports = 1
-            if len(v) < require_ports:
+            if len(value) < require_ports:
                 raise exceptions.PortConfigNotEnough(require_ports)
-        return v
+        return value
 
     def check_port_groups_and_peers(self) -> None:
         topology = self.test_configuration.topology_config.topology
@@ -89,37 +89,35 @@ def check_port_groups_and_peers(self) -> None:
                 if not i:
                     raise exceptions.PortGroupError(group)
 
-    @validator("ports_configuration", always=True)
-    def check_modifier_mode_and_segments(
-        cls, v: "PortConfType", values: Dict[str, Any]
-    ) -> "PortConfType":
-        if "test_configuration" in values:
-            flow_creation_type = values[
+
+    @field_validator("ports_configuration")
+    def check_modifier_mode_and_segments(cls, value: PortConfType, info: ValidationInfo) -> PortConfType:
+        if "test_configuration" in info.data:
+            flow_creation_type = info.data[
                 "test_configuration"
             ].test_execution_config.flow_creation_config.flow_creation_type
-            for port_config in v:
+            for port_config in value:
                 if (
                     not flow_creation_type.is_stream_based
                 ) and port_config.profile.protocol_version.is_l3:
                     raise exceptions.ModifierBasedNotSupportL3()
-        return v
+        return value
+
 
-    @validator("ports_configuration", always=True)
-    def check_port_group(
-        cls, v: "PortConfiguration", values: Dict[str, Any]
-    ) -> "PortConfiguration":
-        if "ports_configuration" in values and "test_configuration" in values:
-            for k, p in values["ports_configuration"].items():
+    @field_validator("ports_configuration")
+    def check_port_group(cls, value: PortConfiguration, info: ValidationInfo) -> PortConfiguration:
+        if "ports_configuration" in info.data and "test_configuration" in info.data:
+            for k, p in info.data["ports_configuration"].items():
                 if (
                     p.port_group == const.PortGroup.UNDEFINED
-                    and not values[
+                    and not info.data[
                         "test_configuration"
                     ].topology_config.topology.is_mesh_topology
                 ):
                     raise exceptions.PortGroupNeeded()
-        return v
+        return value
 
-    @validator("test_types_configuration", always=True)
+    @field_validator("test_types_configuration")
     def check_test_type_enable(
         cls, v: "TestTypesConfiguration"
     ) -> "TestTypesConfiguration":
@@ -134,22 +132,21 @@ def check_test_type_enable(
             raise exceptions.TestTypesError()
         return v
 
-    @validator("test_types_configuration", always=True)
-    def check_result_scope(
-        cls, v: "TestTypesConfiguration", values: Dict[str, Any]
-    ) -> "TestTypesConfiguration":
-        if "test_configuration" not in values:
-            return v
+
+    @field_validator("test_types_configuration")
+    def check_result_scope(cls, value: "TestTypesConfiguration", info: ValidationInfo) -> "TestTypesConfiguration":
+        if "test_configuration" not in info.data:
+            return value
         if (
-            v.throughput_test.enabled
-            and v.throughput_test.rate_iteration_options.result_scope
+            value.throughput_test.enabled
+            and value.throughput_test.rate_iteration_options.result_scope
             == const.RateResultScopeType.PER_SOURCE_PORT
-            and not values[
+            and not info.data[
                 "test_configuration"
             ].test_execution_config.flow_creation_config.flow_creation_type.is_stream_based
         ):
             raise exceptions.ModifierBasedNotSupportPerPortResult()
-        return v
+        return value
 
     @staticmethod
     def count_port_group(

plugin2544/model/m_port_config.py

@@ -7,95 +7,54 @@
     IPv6Address as OriginIPv6Address,
 )
 from typing import Union, Optional
-from pydantic import BaseModel, validator, Field
+from pydantic import BaseModel, field_validator, Field
 from ..utils import constants as const
 from ..utils.field import MacAddress, IPv4Address, IPv6Address, Prefix
 from .m_protocol_segment import ProtocolSegmentProfileConfig
 
 
-class IPAddressProperties(BaseModel):
-    address: Union[IPv4Address, IPv6Address] = IPv4Address("0.0.0.0")
-    routing_prefix: Prefix = Prefix(24)
-    public_address: Union[IPv4Address, IPv6Address] = IPv4Address("0.0.0.0")
-    public_routing_prefix: Prefix = Prefix(24)
-    gateway: Union[IPv4Address, IPv6Address] = IPv4Address("0.0.0.0")
-    remote_loop_address: Union[IPv4Address, IPv6Address] = IPv4Address("0.0.0.0")
+class IPAddressProperties(BaseModel, arbitrary_types_allowed=True):
+    address: Union[IPv4Address, IPv6Address, str] = "0.0.0.0"
+    routing_prefix: Prefix | int = Prefix(24)
+    public_address: Union[IPv4Address, IPv6Address, str] = "0.0.0.0"
+    public_routing_prefix: Prefix | int = Prefix(24)
+    gateway: Union[IPv4Address, IPv6Address, str] = "0.0.0.0"
+    remote_loop_address: Union[IPv4Address, IPv6Address, str] = "0.0.0.0"
     # ip_version: const.IPVersion = const.IPVersion.IPV6
 
     @property
     def network(self) -> Union["IPv4Network", "IPv6Network"]:
         return ip_network(f"{self.address}/{self.routing_prefix}", strict=False)
 
-    @validator(
-        "address",
-        "public_address",
-        "gateway",
-        "remote_loop_address",
-        pre=True,
-        allow_reuse=True,
-    )
+    @field_validator("address", "public_address", "gateway", "remote_loop_address", mode="before")
     def set_address(
-        cls, origin_addr: Union[str, "IPv4Address", "IPv6Address"]
+        cls, value: Union[str, "IPv4Address", "IPv6Address"]
     ) -> Union["IPv4Address", "IPv6Address"]:
-        address = ip_address(origin_addr)
+        address = ip_address(value)
         return (
             IPv4Address(address)
             if isinstance(address, OriginIPv4Address)
             else IPv6Address(address)
         )
 
-    @validator("routing_prefix", "public_routing_prefix", pre=True, allow_reuse=True)
-    def set_prefix(cls, v: int) -> Prefix:
-        return Prefix(v)
+    @field_validator("routing_prefix", "public_routing_prefix", mode="before")
+    def set_prefix(cls, value: int) -> Prefix:
+        return Prefix(value)
 
     @property
-    def dst_addr(self) -> Union["IPv4Address", "IPv6Address"]:
+    def dst_addr(self) -> Union["IPv4Address", "IPv6Address", str]:
         return self.public_address if not self.public_address.is_empty else self.address
 
-
-# class IPV4AddressProperties(BaseModel):
-#     address: IPv4Address = IPv4Address("0.0.0.0")
-#     routing_prefix: Prefix = Prefix(24)
-#     public_address: IPv4Address = IPv4Address("0.0.0.0")
-#     public_routing_prefix: Prefix = Prefix(24)
-#     gateway: IPv4Address = IPv4Address("0.0.0.0")
-#     remote_loop_address: IPv4Address = IPv4Address("0.0.0.0")
-#     ip_version: const.IPVersion = const.IPVersion.IPV4
-
-#     @property
-#     def network(self) -> "IPv4Network":
-#         return IPv4Network(f"{self.address}/{self.routing_prefix}", strict=False)
-
-#     @validator(
-#         "address",
-#         "public_address",
-#         "gateway",
-#         "remote_loop_address",
-#         pre=True,
-#         allow_reuse=True,
-#     )
-#     def set_address(cls, v: Union[str, "IPv4Address"]) -> "IPv4Address":
-#         return IPv4Address(v)
-
-#     @validator("routing_prefix", "public_routing_prefix", pre=True, allow_reuse=True)
-#     def set_prefix(cls, v: int) -> Prefix:
-#         return Prefix(v)
-
-#     @property
-#     def dst_addr(self) -> "IPv4Address":
-#         return self.public_address if not self.public_address.is_empty else self.address
-
-
-class PortConfiguration(BaseModel):
+class PortConfiguration(BaseModel, arbitrary_types_allowed=True):
     port_slot: int
     peer_slot: Optional[int]
     port_group: const.PortGroup
     port_speed_mode: const.PortSpeedStr
     ip_address: Optional[IPAddressProperties]
-    ip_gateway_mac_address: MacAddress
+    ip_gateway_mac_address: MacAddress | str
     reply_arp_requests: bool
     reply_ping_requests: bool
-    remote_loop_mac_address: MacAddress
+    remote_loop_mac_address: MacAddress | str
     inter_frame_gap: float
     speed_reduction_ppm: int = Field(ge=0)
     pause_mode_enabled: bool
@@ -121,12 +80,9 @@ class PortConfiguration(BaseModel):
     _is_tx: bool = True
     _is_rx: bool = True
 
-    class Config:
-        underscore_attrs_are_private = True
-
-    @validator("ip_gateway_mac_address", pre=True)
-    def set_ip_gateway_mac_address(cls, ip_gateway_mac_address: str) -> "MacAddress":
-        return MacAddress(ip_gateway_mac_address)
+    @field_validator("ip_gateway_mac_address", mode="before")
+    def set_ip_gateway_mac_address(cls, value: str) -> "MacAddress":
+        return MacAddress(value)
 
     @property
     def is_tx_port(self) -> bool:

plugin2544/model/m_protocol_segment.py

@@ -1,22 +1,27 @@
 import re
 from enum import Enum
 from random import randint
-from typing import Any, Callable, Dict, Generator, List, Optional
-from pydantic import BaseModel, Field
-from pydantic.class_validators import validator
+from typing import Any, Callable, Dict, Generator, List, Optional, Annotated
+from pydantic import BaseModel, Field, field_validator, ValidationInfo
+from pydantic_core import CoreSchema, core_schema
+from pydantic import GetCoreSchemaHandler, TypeAdapter
 from xoa_driver.enums import ProtocolOption, ModifierAction
 from ..utils.exceptions import ModifierRangeError
 
 class BinaryString(str):
     @classmethod
-    def __get_validators__(cls) -> Generator[Callable, None, None]:
-        yield cls.validate
-
-    @classmethod
-    def validate(cls, v: str) -> "BinaryString":
-        if not re.search("^[01]+$", v):
-            raise ValueError("binary string must zero or one")
-        return cls(v)
+    def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema:
+        return core_schema.no_info_after_validator_function(cls, handler(str))
+    
+    # @classmethod
+    # def __get_validators__(cls) -> Generator[Callable, None, None]:
+    #     yield cls.validate
+
+    # @classmethod
+    # def validate(cls, v: str) -> "BinaryString":
+    #     if not re.search("^[01]+$", v):
+    #         raise ValueError("binary string must zero or one")
+    #     return cls(v)
 
     @property
     def is_all_zero(self) -> bool:
@@ -91,13 +96,17 @@ class SegmentType(Enum):
 
     @property
     def is_raw(self) -> bool:
-        return self.value.lower().startswith("raw")
+        if isinstance(self.value, str):
+            return self.value.lower().startswith("raw")
+        return False
 
     @property
     def raw_length(self) -> int:
         if not self.is_raw:
             return 0
-        return int(self.value.split("_")[-1])
+        if isinstance(self.value, str):
+            return int(self.value.split("_")[-1])
+        return 0
 
     def to_xmp(self) -> "ProtocolOption":
         return ProtocolOption[self.name]
@@ -123,9 +132,6 @@ class ValueRange(BaseModel):
     restart_for_each_port: bool
     _current_count: int = 0  # counter start from 0
 
-    class Config:
-        underscore_attrs_are_private = True
-
     def reset(self) -> None:
         self._current_count = 0
 
@@ -157,21 +163,20 @@ def get_current_value(self) -> int:
 class HWModifier(BaseModel):
     start_value: int
     step_value: int = Field(gt=0)
-    stop_value: int
+    stop_value: Annotated[int, Field(validate_default=True)]
     repeat: int
     offset: int
     action: ModifierActionOption
     mask: str  # hex string as 'FFFF'
     _byte_segment_position: int = 0  # byte position of all header segments
 
-    class Config:
-        underscore_attrs_are_private = True
-
-    @validator('stop_value', pre=True, always=True)
-    def validate_modifier_value(cls, v: int, values: Dict[str, Any]):
-        if (v - values['start_value']) % values['step_value']:
-            raise ModifierRangeError(values['start_value'], v, values['step_value'])
-        return v
+    @field_validator('stop_value', mode="before")
+    def validate_modifier_value(cls, value: int, info: ValidationInfo):
+        start_value = info.data['start_value']
+        step_value = info.data['step_value']
+        if (value - start_value) % step_value:
+            raise ModifierRangeError(start_value, value, step_value)
+        return value
 
 
     def set_byte_segment_position(self, position: int) -> None:
@@ -241,7 +246,7 @@ def hw_modifiers(self) -> Generator["HWModifier", None, None]:
     def value_ranges(self) -> Generator["ValueRange", None, None]:
         return (f.value_range for f in self.fields if f.value_range)
 
-    @validator("checksum_offset")
+    @field_validator("checksum_offset")
     def is_digit(cls, value: int) -> int:
         if value and not isinstance(value, int):
             raise ValueError("checksum offset must digit")