Fix compatibility with pydantic v2

Kircheneer · Kircheneer · commit 4596499259fd · 2024-01-22T08:47:19.000+01:00
diff --git a/diffsync/__init__.py b/diffsync/__init__.py
@@ -106,24 +106,24 @@ class DiffSyncModel(BaseModel):
     """Message, if any, associated with the create/update/delete status value."""
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def __init_subclass__(cls) -> None:
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
         """Validate that the various class attribute declarations correspond to actual instance fields.
 
         Called automatically on subclass declaration.
         """
-        variables = cls.__fields__.keys()
         # Make sure that any field referenced by name actually exists on the model
         for attr in cls._identifiers:
-            if attr not in variables and not hasattr(cls, attr):
+            if attr not in cls.model_fields and not hasattr(cls, attr):
                 raise AttributeError(f"_identifiers {cls._identifiers} references missing or un-annotated attr {attr}")
         for attr in cls._shortname:
-            if attr not in variables:
+            if attr not in cls.model_fields:
                 raise AttributeError(f"_shortname {cls._shortname} references missing or un-annotated attr {attr}")
         for attr in cls._attributes:
-            if attr not in variables:
+            if attr not in cls.model_fields:
                 raise AttributeError(f"_attributes {cls._attributes} references missing or un-annotated attr {attr}")
         for attr in cls._children.values():
-            if attr not in variables:
+            if attr not in cls.model_fields:
                 raise AttributeError(f"_children {cls._children} references missing or un-annotated attr {attr}")
 
         # Any given field can only be in one of (_identifiers, _attributes, _children)
@@ -147,15 +147,15 @@ def dict(self, **kwargs: Any) -> Dict:
         """Convert this DiffSyncModel to a dict, excluding the diffsync field by default as it is not serializable."""
         if "exclude" not in kwargs:
             kwargs["exclude"] = {"diffsync"}
-        return super().dict(**kwargs)
+        return super().model_dump(**kwargs)
 
     def json(self, **kwargs: Any) -> StrType:
         """Convert this DiffSyncModel to a JSON string, excluding the diffsync field by default as it is not serializable."""
         if "exclude" not in kwargs:
             kwargs["exclude"] = {"diffsync"}
         if "exclude_defaults" not in kwargs:
             kwargs["exclude_defaults"] = True
-        return super().json(**kwargs)
+        return super().model_dump_json(**kwargs)
 
     def str(self, include_children: bool = True, indent: int = 0) -> StrType:
         """Build a detailed string representation of this DiffSyncModel and optionally its children."""
@@ -855,4 +855,4 @@ def count(self, model: Union[StrType, "DiffSyncModel", Type["DiffSyncModel"], No
 
 
 # DiffSyncModel references DiffSync and DiffSync references DiffSyncModel. Break the typing loop:
-DiffSyncModel.update_forward_refs()
+DiffSyncModel.model_rebuild()
diff --git a/examples/01-multiple-data-sources/models.py b/examples/01-multiple-data-sources/models.py
@@ -40,8 +40,8 @@ class Device(DiffSyncModel):
     _children = {"interface": "interfaces"}
 
     name: str
-    site_name: Optional[str]  # note that this attribute is NOT included in _attributes
-    role: Optional[str]  # note that this attribute is NOT included in _attributes
+    site_name: Optional[str] = None  # note that this attribute is NOT included in _attributes
+    role: Optional[str] = None  # note that this attribute is NOT included in _attributes
     interfaces: List = []
 
 
@@ -56,4 +56,4 @@ class Interface(DiffSyncModel):
     name: str
     device_name: str
 
-    description: Optional[str]
+    description: Optional[str] = None
diff --git a/examples/03-remote-system/models.py b/examples/03-remote-system/models.py
@@ -32,4 +32,4 @@ class Country(DiffSyncModel):
     slug: str
     name: str
     region: str
-    population: Optional[int]
+    population: Optional[int] = 0
diff --git a/examples/04-get-update-instantiate/models.py b/examples/04-get-update-instantiate/models.py
@@ -38,8 +38,8 @@ class Device(DiffSyncModel):
     _children = {"interface": "interfaces", "site": "sites"}
 
     name: str
-    site_name: Optional[str]  # note that this attribute is NOT included in _attributes
-    role: Optional[str]  # note that this attribute is NOT included in _attributes
+    site_name: Optional[str] = None  # note that this attribute is NOT included in _attributes
+    role: Optional[str] = None  # note that this attribute is NOT included in _attributes
     interfaces: List = []
     sites: List = []
 
@@ -55,4 +55,4 @@ class Interface(DiffSyncModel):
     name: str
     device_name: str
 
-    description: Optional[str]
+    description: Optional[str] = None
diff --git a/examples/05-nautobot-peeringdb/models.py b/examples/05-nautobot-peeringdb/models.py
@@ -21,8 +21,8 @@ class RegionModel(DiffSyncModel):
     # Data type declarations for all identifiers and attributes
     name: str
     slug: str
-    description: Optional[str]
-    parent_name: Optional[str]  # may be None
+    description: Optional[str] = None
+    parent_name: Optional[str] = None
     sites: List = []
 
     # Not in _attributes or _identifiers, hence not included in diff calculations
@@ -49,10 +49,10 @@ class SiteModel(DiffSyncModel):
     name: str
     slug: str
     status_slug: str
-    region_name: Optional[str]  # may be None
-    description: Optional[str]
-    latitude: Optional[float]
-    longitude: Optional[float]
+    region_name: Optional[str] = None
+    description: Optional[str] = None
+    latitude: Optional[float] = None
+    longitude: Optional[float] = None
 
     # Not in _attributes or _identifiers, hence not included in diff calculations
     pk: Optional[Union[UUID, int]]
diff --git a/examples/06-ip-prefixes/models.py b/examples/06-ip-prefixes/models.py
@@ -11,6 +11,6 @@ class Prefix(DiffSyncModel):
     _attributes = ("vrf", "vlan_id", "tenant")
 
     prefix: str
-    vrf: Optional[str]
-    vlan_id: Optional[int]
-    tenant: Optional[str]
+    vrf: Optional[str] = None
+    vlan_id: Optional[int] = None
+    tenant: Optional[str] = None
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
@@ -115,7 +115,7 @@ class Device(DiffSyncModel):
     _children = {"interface": "interfaces"}
 
     name: str
-    site_name: Optional[str]  # note this is not included in _attributes
+    site_name: Optional[str] = None  # note this is not included in _attributes
     role: str
     interfaces: List = []
 
@@ -143,7 +143,7 @@ class Interface(DiffSyncModel):
     name: str
 
     interface_type: str = "ethernet"
-    description: Optional[str]
+    description: Optional[str] = None
 
 
 @pytest.fixture
diff --git a/tests/unit/test_diffsync_model.py b/tests/unit/test_diffsync_model.py
@@ -107,7 +107,7 @@ def test_diffsync_model_dict_with_data(make_interface):
 def test_diffsync_model_json_with_data(make_interface):
     intf = make_interface()
     # json() omits default values for brevity
-    assert intf.json() == '{"device_name": "device1", "name": "eth0"}'
+    assert intf.json() == '{"device_name":"device1","name":"eth0"}'
 
 
 def test_diffsync_model_str_with_data(make_interface):
@@ -182,7 +182,7 @@ def test_diffsync_model_json_with_children(generic_diffsync, make_site, make_dev
     generic_diffsync.add(site1)
     generic_diffsync.add(device1)
 
-    assert site1.json() == '{"name": "site1", "devices": ["device1"]}'
+    assert site1.json() == '{"name":"site1","devices":["device1"]}'
 
 
 def test_diffsync_model_str_with_children(generic_diffsync, make_site, make_device, make_interface):
@@ -294,7 +294,7 @@ class BadAttributes(DiffSyncModel):
 
             name: str
             # Note that short_name doesn't have a type annotation - making sure this works too
-            short_name = "short_name"
+            short_name: str = "short_name"
 
     assert "_attributes" in str(excinfo.value)
     assert "my_attr" in str(excinfo.value)
@@ -310,7 +310,7 @@ class BadChildren(DiffSyncModel):
             _children = {"device": "devices"}
 
             name: str
-            short_name = "short_name"
+            short_name: str = "short_name"
             my_attr: int = 0
 
     assert "_children" in str(excinfo.value)
