Migrate to Pydantic v2

Author: istride

src/rpft/parsers/common/model_inference.py

@@ -1,7 +1,7 @@
 from collections import defaultdict
 from typing import List, ForwardRef, _eval_type
 from pydoc import locate
-from pydantic.v1 import create_model
+from pydantic import create_model
 
 from rpft.parsers.common.rowparser import (
     ParserModel,

src/rpft/parsers/common/rowparser.py

@@ -3,7 +3,7 @@
 from collections.abc import Iterable
 from typing import List
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from rpft.parsers.common.cellparser import CellParser
 
@@ -97,9 +97,7 @@ def is_basic_instance(value):
 
 
 def is_default_value(model_instance, field, field_value):
-    # Note: In pydantic V2, __fields__ will become model_fields
-    if field_value == type(model_instance).__fields__[field].get_default():
-        return True
+    return field_value == type(model_instance).model_fields[field].default
 
 
 def str_to_bool(string):
@@ -137,12 +135,12 @@ def try_assign_as_kwarg(self, field, key, value, model):
         # model, assign value to field[key] (which represents the field in the model)
         if is_list_instance(value) and len(value) == 2 and type(value[0]) is str:
             first_entry_as_key = model.header_name_to_field_name(value[0])
-            if first_entry_as_key in model.__fields__:
+            if first_entry_as_key in model.model_fields:
                 self.assign_value(
                     field[key],
                     first_entry_as_key,
                     value[1],
-                    model.__fields__[first_entry_as_key].outer_type_,
+                    model.model_fields[first_entry_as_key].annotation,
                 )
                 return True
         return False
@@ -164,7 +162,7 @@ def assign_value(self, field, key, value, model):
             # Get the list of keys that are available for the target model
             # Note: The fields have a well defined ordering.
             # See https://pydantic-docs.helpmanual.io/usage/models/#field-ordering
-            model_fields = list(model.__fields__.keys())
+            model_fields = list(model.model_fields.keys())
 
             if type(value) is not list:
                 # It could be that an object is specified via a single element.
@@ -193,7 +191,7 @@ def assign_value(self, field, key, value, model):
                         field[key],
                         entry_key,
                         entry,
-                        model.__fields__[entry_key].outer_type_,
+                        model.model_fields[entry_key].annotation,
                     )
         elif is_basic_dict_type(model):
             field[key] = {}
@@ -281,12 +279,9 @@ def find_entry(self, model, output_field, field_path):
         else:
             assert is_parser_model_type(model)
             key = model.header_name_to_field_name(field_name)
-            if key not in model.__fields__:
+            if key not in model.model_fields:
                 raise ValueError(f"Field {key} doesn't exist in target type {model}.")
-            child_model = model.__fields__[key].outer_type_
-            # TODO: how does ModelField.outer_type_ and ModelField.type_
-            # deal with nested lists, e.g. List[List[str]]?
-            # Write test cases and fix code.
+            child_model = model.model_fields[key].annotation
 
             if key not in output_field:
                 # Create a new entry for this, if necessary

src/rpft/parsers/creation/campaigneventrowmodel.py

@@ -1,5 +1,6 @@
+from pydantic import field_validator
+
 from rpft.parsers.common.rowparser import ParserModel
-from pydantic.v1 import validator
 
 
 class CampaignEventRowModel(ParserModel):
@@ -14,13 +15,15 @@ class CampaignEventRowModel(ParserModel):
     flow: str = ""
     base_language: str = ""
 
-    @validator("unit")
+    @field_validator("unit")
+    @classmethod
     def validate_unit(cls, v):
         if v not in ["M", "H", "D", "W"]:
             raise ValueError("unit must be M (minute), H (hour), D (day) or W (week)")
         return v
 
-    @validator("start_mode")
+    @field_validator("start_mode")
+    @classmethod
     def validate_start_mode(cls, v):
         if v not in ["I", "S", "P"]:
             raise ValueError(
@@ -29,7 +32,8 @@ def validate_start_mode(cls, v):
             )
         return v
 
-    @validator("event_type")
+    @field_validator("event_type")
+    @classmethod
     def validate_event_type(cls, v):
         if v not in ["M", "F"]:
             raise ValueError("event_type must be F (flow) or M (message)")

src/rpft/parsers/creation/contentindexparser.py

@@ -36,7 +36,7 @@ def __init__(self, rows, row_model):
     def to_dict(self):
         return {
             "model": self.row_model.__name__,
-            "rows": [content.dict() for content in self.rows.values()],
+            "rows": [content.model_dump() for content in self.rows.values()],
         }
 
 

src/rpft/parsers/creation/flowrowmodel.py

@@ -1,3 +1,5 @@
+from pydantic import ConfigDict
+
 from rpft.parsers.common.rowparser import ParserModel
 from rpft.parsers.creation.models import Condition
 
@@ -45,6 +47,8 @@ class WhatsAppTemplating(ParserModel):
 
 
 class Edge(ParserModel):
+    model_config = ConfigDict(coerce_numbers_to_str=True)
+
     from_: str = ""
     condition: Condition = Condition()
 
@@ -65,6 +69,8 @@ def header_name_to_field_name_with_context(header, row):
 
 
 class FlowRowModel(ParserModel):
+    model_config = ConfigDict(coerce_numbers_to_str=True)
+
     row_id: str = ""
     type: str
     edges: list[Edge]

src/rpft/parsers/creation/triggerrowmodel.py

@@ -1,4 +1,4 @@
-from pydantic.v1 import validator
+from pydantic import field_validator, model_validator
 
 from rpft.parsers.common.rowparser import ParserModel
 
@@ -12,7 +12,8 @@ class TriggerRowModel(ParserModel):
     channel: str = ""
     match_type: str = ""
 
-    @validator("type")
+    @field_validator("type")
+    @classmethod
     def validate_type(cls, v):
         if v not in ["K", "C", "M", "T"]:
             raise ValueError(
@@ -21,10 +22,11 @@ def validate_type(cls, v):
             )
         return v
 
-    @validator("match_type")
-    def validate_match_type(cls, v, values):
-        if values["type"] == "K" and v not in ["F", "O", ""]:
+    @model_validator(mode="after")
+    def validate_match_type(self):
+        if self.type == "K" and self.match_type not in ["F", "O", ""]:
             raise ValueError(
                 'match_type must be "F" (starts with) or "O" (only) if type is "K".'
             )
-        return v
+
+        return self

tests/test_differentways.py

@@ -124,11 +124,9 @@ def test_different_ways(self):
         for inp in inputs:
             out = self.parser.parse_row(inp)  # We get an instance of the model
             outputs.append(out)
-            # Note: we can also serialize via out.json(indent=4) for printing
-            # or out.dict()
 
         for out in outputs:
-            self.assertEqual(out.dict(), output_instance)
+            self.assertEqual(out.model_dump(), output_instance)
 
     def test_single_kwarg(self):
         output_single_kwarg = self.parser.parse_row(input_single_kwarg)

tests/test_flowparser_reverse.py

@@ -10,8 +10,8 @@ def setUp(self) -> None:
         pass
 
     def compare_models_leniently(self, model, model_exp):
-        model_dict = model.dict()
-        model_exp_dict = model_exp.dict()
+        model_dict = model.model_dump()
+        model_exp_dict = model_exp.model_dump()
         for edge, edge_exp in zip(model_dict["edges"], model_exp_dict["edges"]):
             if (
                 edge["condition"]["variable"] == "@input.text"

tests/test_model_inference.py

@@ -1,6 +1,6 @@
 from typing import List
 import unittest
-from pydantic.v1 import create_model, BaseModel
+from pydantic import create_model, BaseModel
 
 from rpft.parsers.common.model_inference import (
     get_value_for_type,
@@ -60,7 +60,7 @@ def test_parse_header_annotations(self):
         self.assertEqual(parse_header_annotations("field:int=5"), (int, 5))
 
     def compare_models(self, model1, model2, **kwargs):
-        self.assertEqual(model1(**kwargs).dict(), model2(**kwargs).dict())
+        self.assertEqual(model1(**kwargs).model_dump(), model2(**kwargs).model_dump())
 
     def test_model_from_headers(self):
         self.compare_models(

tests/test_to_row_model.py

@@ -44,8 +44,8 @@ def compare_row_models_without_uuid(self, row_models1, row_models2):
         self.maxDiff = None
         self.assertEqual(len(row_models1), len(row_models2))
         for model1, model2 in zip(row_models1, row_models2):
-            data1 = model1.dict()
-            data2 = model2.dict()
+            data1 = model1.model_dump()
+            data2 = model2.model_dump()
             if not data1["node_uuid"] or not data2["node_uuid"]:
                 # If one of them is blank, skip the comparison
                 data1.pop("node_uuid")

tests/test_unparse.py

@@ -15,7 +15,7 @@ class ModelWithStuff(ParserModel):
 
 class MainModel(ParserModel):
     str_field: str = ""
-    model_optional: Optional[ModelWithStuff]
+    model_optional: Optional[ModelWithStuff] = None
     model_default: ModelWithStuff = ModelWithStuff()
     model_list: List[ModelWithStuff] = []