Adopt Pydantic 2.0

circuit_maintenance_parser/__init__.py

@@ -95,7 +95,6 @@ def get_provider_class(provider_name: str) -> Type[GenericProvider]:
 
 def get_provider_class_from_sender(email_sender: str) -> Type[GenericProvider]:
     """Returns the notification parser class for an email sender address."""
-
     for provider_parser in SUPPORTED_PROVIDERS:
         if provider_parser.get_default_organizer() == email_sender:
             break

circuit_maintenance_parser/data.py

@@ -3,7 +3,7 @@
 from typing import List, NamedTuple, Optional, Type, Set
 
 import email
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel
 from circuit_maintenance_parser.constants import EMAIL_HEADER_SUBJECT, EMAIL_HEADER_DATE
 
 
@@ -18,7 +18,7 @@ class DataPart(NamedTuple):
     content: bytes
 
 
-class NotificationData(BaseModel, extra=Extra.forbid):
+class NotificationData(BaseModel, extra="forbid"):
     """Base class for Notification Data types."""
 
     data_parts: List[DataPart] = []

circuit_maintenance_parser/output.py

@@ -8,7 +8,7 @@
 
 from typing import List
 
-from pydantic import BaseModel, validator, StrictStr, StrictInt, Extra, PrivateAttr
+from pydantic import field_validator, BaseModel, StrictStr, StrictInt, PrivateAttr
 
 
 class Impact(str, Enum):
@@ -52,7 +52,7 @@ class Status(str, Enum):
     NO_CHANGE = "NO-CHANGE"
 
 
-class CircuitImpact(BaseModel, extra=Extra.forbid):
+class CircuitImpact(BaseModel, extra="forbid"):
     """CircuitImpact class.
 
     Each Circuit Maintenance can contain multiple affected circuits, and each one can have a different level of impact.
@@ -73,23 +73,31 @@ class CircuitImpact(BaseModel, extra=Extra.forbid):
         ... )
         Traceback (most recent call last):
         ...
-        pydantic.error_wrappers.ValidationError: 1 validation error for CircuitImpact
+        pydantic_core._pydantic_core.ValidationError: 1 validation error for CircuitImpact
         impact
-          value is not a valid enumeration member; permitted: 'NO-IMPACT', 'REDUCED-REDUNDANCY', 'DEGRADED', 'OUTAGE' (type=type_error.enum; enum_values=[<Impact.NO_IMPACT: 'NO-IMPACT'>, <Impact.REDUCED_REDUNDANCY: 'REDUCED-REDUNDANCY'>, <Impact.DEGRADED: 'DEGRADED'>, <Impact.OUTAGE: 'OUTAGE'>])
+          Input should be 'NO-IMPACT', 'REDUCED-REDUNDANCY', 'DEGRADED' or 'OUTAGE' [type=enum, input_value='wrong impact', input_type=str]
     """
 
     circuit_id: StrictStr
     # Optional Attributes
     impact: Impact = Impact.OUTAGE
 
     # pylint: disable=no-self-argument
-    @validator("impact")
+    @field_validator("impact")
+    @classmethod
     def validate_impact_type(cls, value):
         """Validate Impact type."""
         if value not in Impact:
             raise ValueError("Not a valid impact type")
         return value
 
+    def to_json(self):
+        """Return a JSON serializable dict."""
+        return {
+            "circuit_id": self.circuit_id,
+            "impact": self.impact.value,
+        }
+
 
 class Metadata(BaseModel):
     """Metadata class to provide context about the Maintenance object."""
@@ -100,7 +108,7 @@ class Metadata(BaseModel):
     generated_by_llm: bool = False
 
 
-class Maintenance(BaseModel, extra=Extra.forbid):
+class Maintenance(BaseModel, extra="forbid"):
     """Maintenance class.
 
     Mandatory attributes:
@@ -164,34 +172,44 @@ class Maintenance(BaseModel, extra=Extra.forbid):
 
     def __init__(self, **data):
         """Initialize the Maintenance object."""
-        self._metadata = data.pop("_metadata")
+        metadata = data.pop("_metadata")
         super().__init__(**data)
+        self._metadata = metadata
 
-    # pylint: disable=no-self-argument
-    @validator("status")
+    @field_validator("status")
+    @classmethod
     def validate_status_type(cls, value):
         """Validate Status type."""
         if value not in Status:
             raise ValueError("Not a valid status type")
         return value
 
-    @validator("provider", "account", "maintenance_id", "organizer")
+    @field_validator("provider", "account", "maintenance_id", "organizer")
+    @classmethod
     def validate_empty_strings(cls, value):
         """Validate emptry strings."""
         if value in ["", "None"]:
             raise ValueError("String is empty or 'None'")
         return value
 
-    @validator("circuits")
+    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
+    @field_validator("circuits")
+    @classmethod
     def validate_empty_circuits(cls, value, values):
         """Validate non-cancel notifications have a populated circuit list."""
+        values = values.data
         if len(value) < 1 and str(values["status"]) in ("CANCELLED", "COMPLETED"):
             raise ValueError("At least one circuit has to be included in the maintenance")
         return value
 
-    @validator("end")
+    # TODO[pydantic]: We couldn't refactor the `validator`, please replace it by `field_validator` manually.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-validators for more information.
+    @field_validator("end")
+    @classmethod
     def validate_end_time(cls, end, values):
         """Validate that End time happens after Start time."""
+        values = values.data
         if "start" not in values:
             raise ValueError("Start time is a mandatory attribute.")
         start = values["start"]
@@ -209,6 +227,6 @@ def to_json(self) -> str:
         return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=2)
 
     @property
-    def metadata(self):
+    def metadata(self) -> Metadata:
         """Get Maintenance Metadata."""
         return self._metadata

circuit_maintenance_parser/parser.py

@@ -12,7 +12,7 @@
 import bs4  # type: ignore
 from bs4.element import ResultSet  # type: ignore
 
-from pydantic import BaseModel
+from pydantic import BaseModel, PrivateAttr
 from icalendar import Calendar  # type: ignore
 
 from circuit_maintenance_parser.errors import ParserError
@@ -34,15 +34,15 @@ class Parser(BaseModel):
     """
 
     # _data_types are used to match the Parser to to each type of DataPart
-    _data_types = ["text/plain", "plain"]
+    _data_types = PrivateAttr(["text/plain", "plain"])
 
     # TODO: move it to where it is used, Cogent parser
     _geolocator = Geolocator()
 
     @classmethod
     def get_data_types(cls) -> List[str]:
         """Return the expected data type."""
-        return cls._data_types
+        return cls._data_types.get_default()
 
     @classmethod
     def get_name(cls) -> str:
@@ -92,7 +92,7 @@ class ICal(Parser):
     Reference: https://tools.ietf.org/html/draft-gunter-calext-maintenance-notifications-00
     """
 
-    _data_types = ["text/calendar", "ical", "icalendar"]
+    _data_types = PrivateAttr(["text/calendar", "ical", "icalendar"])
 
     def parser_hook(self, raw: bytes, content_type: str):
         """Execute parsing."""
@@ -164,7 +164,7 @@ def parse_ical(gcal: Calendar) -> List[Dict]:
 class Html(Parser):
     """Html parser."""
 
-    _data_types = ["text/html", "html"]
+    _data_types = PrivateAttr(["text/html", "html"])
 
     @staticmethod
     def remove_hex_characters(string):
@@ -201,7 +201,11 @@ def clean_line(line):
 class EmailDateParser(Parser):
     """Parser for Email Date."""
 
-    _data_types = [EMAIL_HEADER_DATE]
+    _data_types = PrivateAttr(
+        [
+            EMAIL_HEADER_DATE,
+        ]
+    )
 
     def parser_hook(self, raw: bytes, content_type: str):
         """Execute parsing."""
@@ -214,7 +218,11 @@ def parser_hook(self, raw: bytes, content_type: str):
 class EmailSubjectParser(Parser):
     """Parse data from subject or email."""
 
-    _data_types = [EMAIL_HEADER_SUBJECT]
+    _data_types = PrivateAttr(
+        [
+            EMAIL_HEADER_SUBJECT,
+        ]
+    )
 
     def parser_hook(self, raw: bytes, content_type: str):
         """Execute parsing."""
@@ -236,7 +244,7 @@ def bytes_to_string(string):
 class Csv(Parser):
     """Csv parser."""
 
-    _data_types = ["application/csv", "text/csv", "application/octet-stream"]
+    _data_types = PrivateAttr(["application/csv", "text/csv", "application/octet-stream"])
 
     def parser_hook(self, raw: bytes, content_type: str):
         """Execute parsing."""
@@ -255,7 +263,11 @@ def parse_csv(raw: bytes) -> List[Dict]:
 class Text(Parser):
     """Text parser."""
 
-    _data_types = ["text/plain"]
+    _data_types = PrivateAttr(
+        [
+            "text/plain",
+        ]
+    )
 
     def parser_hook(self, raw: bytes, content_type: str):
         """Execute parsing."""
@@ -278,7 +290,7 @@ def parse_text(self, text) -> List[Dict]:
 class LLM(Parser):
     """LLM parser."""
 
-    _data_types = ["text/html", "html", "text/plain"]
+    _data_types = PrivateAttr(["text/html", "html", "text/plain"])
 
     _llm_question = """Please, could you extract a JSON form without any other comment,
     with the following JSON schema (timestamps in EPOCH and taking into account the GMT offset):

circuit_maintenance_parser/processor.py

@@ -5,8 +5,7 @@
 
 from typing import Iterable, Type, Dict, List
 
-from pydantic import BaseModel, Extra
-from pydantic.error_wrappers import ValidationError
+from pydantic import BaseModel, ValidationError
 
 from circuit_maintenance_parser.output import Maintenance, Metadata
 from circuit_maintenance_parser.data import NotificationData
@@ -17,7 +16,7 @@
 logger = logging.getLogger(__name__)
 
 
-class GenericProcessor(BaseModel, extra=Extra.forbid):
+class GenericProcessor(BaseModel, extra="forbid"):
     """Base class for the Processors.
 
     Attributes: