Skip to content

Commit

Permalink
Adopt Pydantic 2.0 (#264)
Browse files Browse the repository at this point in the history
* Bump mypy

* Adjust code for pydanticv2

* Remove pydantic comments

* add todo for nestedencoder

* constraint mypy
  • Loading branch information
chadell authored Feb 14, 2024
1 parent 57ad93d commit 325a834
Show file tree
Hide file tree
Showing 11 changed files with 691 additions and 519 deletions.
1 change: 0 additions & 1 deletion circuit_maintenance_parser/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions circuit_maintenance_parser/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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] = []
Expand Down
40 changes: 27 additions & 13 deletions circuit_maintenance_parser/output.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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.
Expand All @@ -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."""
Expand All @@ -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:
Expand Down Expand Up @@ -164,34 +172,40 @@ 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")
@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")
@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"]
Expand All @@ -209,6 +223,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
32 changes: 22 additions & 10 deletions circuit_maintenance_parser/parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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."""
Expand All @@ -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."""
Expand All @@ -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."""
Expand All @@ -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."""
Expand All @@ -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):
Expand Down
5 changes: 2 additions & 3 deletions circuit_maintenance_parser/processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down
Loading

0 comments on commit 325a834

Please sign in to comment.