Upgrade to Pydantic v2

pyproject.toml

@@ -25,7 +25,8 @@ krr = "robusta_krr.main:run"
 [tool.poetry.dependencies]
 python = ">=3.10,<=3.12.9"
 typer = { extras = ["all"], version = "^0.7.0" }
-pydantic = "^1.10.7"
+pydantic = ">=2.0,<3.0"
+pydantic-settings = ">=2.0,<3.0"
 kubernetes = "^26.1.0"
 prometheus-api-client = "0.5.3"
 numpy = ">=1.26.4,<1.27.0"
@@ -35,7 +36,7 @@ slack-sdk = "^3.21.3"
 pandas = "2.2.2"
 requests = ">2.32.4"
 pyyaml = "6.0.1"
-typing-extensions = "4.6.0"
+typing-extensions = ">=4.6.1"
 idna = "3.7"
 urllib3 = "^2.6.2"
 setuptools = "^80.9.0"

requirements.txt

@@ -29,7 +29,8 @@ prometheus-api-client==0.5.3 ; python_version >= "3.10" and python_full_version
 prometrix==0.2.11; python_version >= "3.10" and python_full_version < "3.13"
 pyasn1-modules==0.4.2 ; python_version >= "3.10" and python_full_version < "3.13"
 pyasn1==0.6.2 ; python_version >= "3.10" and python_full_version < "3.13"
-pydantic==1.10.15 ; python_version >= "3.10" and python_full_version < "3.13"
+pydantic>=2.0,<3.0 ; python_version >= "3.10" and python_full_version < "3.13"
+pydantic-settings>=2.0,<3.0 ; python_version >= "3.10" and python_full_version < "3.13"
 pygments==2.17.2 ; python_version >= "3.10" and python_full_version < "3.13"
 pyparsing==3.1.2 ; python_version >= "3.10" and python_full_version < "3.13"
 python-dateutil==2.9.0.post0 ; python_version >= "3.10" and python_full_version < "3.13"
@@ -47,7 +48,7 @@ six==1.16.0 ; python_version >= "3.10" and python_full_version < "3.13"
 slack-sdk==3.27.1 ; python_version >= "3.10" and python_full_version < "3.13"
 tenacity==9.0.0 ; python_version >= "3.10" and python_full_version < "3.13"
 typer==0.7.0 ; python_version >= "3.10" and python_full_version < "3.13"
-typing-extensions==4.6.0 ; python_version >= "3.10" and python_full_version < "3.13"
+typing-extensions>=4.6.1 ; python_version >= "3.10" and python_full_version < "3.13"
 tzdata==2024.1 ; python_version >= "3.10" and python_full_version < "3.13"
 tzlocal==5.2 ; python_version >= "3.10" and python_full_version < "3.13"
 urllib3==2.6.2 ; python_version >= "3.10" and python_full_version < "3.13"

robusta_krr/core/models/allocations.py

@@ -5,6 +5,7 @@
 from typing import Literal, Optional, TypeVar, Union
 
 import pydantic as pd
+from pydantic import field_validator
 from kubernetes.client.models import V1Container
 
 from robusta_krr.utils import resource_units
@@ -52,7 +53,7 @@ def format_diff(allocated, recommended, selector, multiplier=1, colored=False) -
 class ResourceAllocations(pd.BaseModel):
     requests: dict[ResourceType, RecommendationValue]
     limits: dict[ResourceType, RecommendationValue]
-    info: dict[ResourceType, Optional[str]] = {}
+    info: dict[ResourceType, Optional[str]] = pd.Field(default_factory=dict)
 
     @staticmethod
     def __parse_resource_value(value: RecommendationValueRaw) -> RecommendationValue:
@@ -67,7 +68,8 @@ def __parse_resource_value(value: RecommendationValueRaw) -> RecommendationValue
 
         return float(value)
 
-    @pd.validator("requests", "limits", pre=True)
+    @field_validator("requests", "limits", mode="before")
+    @classmethod
     def validate_requests(
         cls, value: dict[ResourceType, RecommendationValueRaw]
     ) -> dict[ResourceType, RecommendationValue]:

robusta_krr/core/models/config.py

@@ -5,6 +5,8 @@
 from typing import Any, Literal, Optional, Union
 
 import pydantic as pd
+from pydantic import field_validator
+from pydantic_settings import BaseSettings
 from kubernetes import config
 from kubernetes.config.config_exception import ConfigException
 from rich.console import Console
@@ -17,7 +19,7 @@
 logger = logging.getLogger("krr")
 
 
-class Config(pd.BaseSettings):
+class Config(BaseSettings):
     quiet: bool = pd.Field(False)
     verbose: bool = pd.Field(False)
 
@@ -98,7 +100,8 @@ def __init__(self, **kwargs: Any) -> None:
     def Formatter(self) -> formatters.FormatterFunc:
         return formatters.find(self.format)
 
-    @pd.validator("prometheus_url")
+    @field_validator("prometheus_url")
+    @classmethod
     def validate_prometheus_url(cls, v: Optional[str]):
         if v is None:
             return None
@@ -110,14 +113,16 @@ def validate_prometheus_url(cls, v: Optional[str]):
 
         return v
 
-    @pd.validator("prometheus_other_headers", pre=True)
+    @field_validator("prometheus_other_headers", mode="before")
+    @classmethod
     def validate_prometheus_other_headers(cls, headers: Union[list[str], dict[str, str]]) -> dict[str, str]:
         if isinstance(headers, dict):
             return headers
 
         return {k.strip().lower(): v.strip() for k, v in [header.split(":") for header in headers]}
 
-    @pd.validator("namespaces")
+    @field_validator("namespaces")
+    @classmethod
     def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]:
         if v == []:
             return "*"
@@ -129,7 +134,8 @@ def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[st
 
         return [val.lower() for val in v]
 
-    @pd.validator("resources", pre=True)
+    @field_validator("resources", mode="before")
+    @classmethod
     def validate_resources(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]:
         if v == []:
             return "*"
@@ -138,7 +144,8 @@ def validate_resources(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str
         # So this will preserve the big and small letters of the resource
         return [next(r for r in KindLiteral.__args__ if r.lower() == val.lower()) for val in v]
 
-    @pd.validator("job_grouping_labels", pre=True)
+    @field_validator("job_grouping_labels", mode="before")
+    @classmethod
     def validate_job_grouping_labels(cls, v: Union[list[str], str, None]) -> Union[list[str], None]:
         if v is None:
             return None
@@ -152,12 +159,14 @@ def create_strategy(self) -> AnyStrategy:
         StrategySettingsType = StrategyType.get_settings_type()
         return StrategyType(StrategySettingsType(**self.other_args))  # type: ignore
 
-    @pd.validator("strategy")
+    @field_validator("strategy")
+    @classmethod
     def validate_strategy(cls, v: str) -> str:
         BaseStrategy.find(v)  # NOTE: raises if strategy is not found
         return v
 
-    @pd.validator("format")
+    @field_validator("format")
+    @classmethod
     def validate_format(cls, v: str) -> str:
         formatters.find(v)  # NOTE: raises if strategy is not found
         return v
@@ -213,7 +222,9 @@ def get_config() -> Optional[Config]:
 
 # NOTE: This class is just a proxy for _config.
 # Import settings from this module and use it like it is just a config object.
-class _Settings(Config):  # Config here is used for type checking
+class _Settings:
+    """Proxy that delegates attribute access to the global _config instance."""
+
     def __init__(self) -> None:
         pass
 

robusta_krr/core/models/objects.py

@@ -45,7 +45,7 @@ class K8sObjectData(pd.BaseModel):
     namespace: str
     kind: KindLiteral
     allocations: ResourceAllocations
-    warnings: set[PodWarning] = set()
+    warnings: set[PodWarning] = pd.Field(default_factory=set)
     labels: Optional[dict[str, str]]
     annotations: Optional[dict[str, str]]
 

robusta_krr/core/models/result.py

@@ -69,8 +69,7 @@ class Result(pd.BaseModel):
     clusterSummary: dict[str, Any] = {}
     config: Optional[Config] = pd.Field(default_factory=Config.get_config)
 
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def model_post_init(self, __context: Any) -> None:
         self.score = self.__calculate_score()
 
     def format(self, formatter: Union[formatters.FormatterFunc, str]) -> Any:

robusta_krr/core/runner.py

@@ -478,7 +478,7 @@ async def _collect_result(self) -> Result:
             description=f"[b]{self._strategy.display_name.title()} Strategy[/b]\n\n{self._strategy.description}",
             strategy=StrategyData(
                 name=str(self._strategy).lower(),
-                settings=self._strategy.settings.dict(),
+                settings=self._strategy.settings.model_dump(),
             ),
             clusterSummary=cluster_summary
         )

robusta_krr/formatters/pprint.py

@@ -6,4 +6,4 @@
 
 @formatters.register()
 def pprint(result: Result) -> str:
-    return pformat(result.dict())
+    return pformat(result.model_dump())

robusta_krr/formatters/yaml.py

@@ -8,4 +8,4 @@
 
 @formatters.register()
 def yaml(result: Result) -> str:
-    return yaml_module.dump(json.loads(result.json()), sort_keys=False)
+    return yaml_module.dump(json.loads(result.model_dump_json()), sort_keys=False)

robusta_krr/main.py

@@ -427,14 +427,14 @@ def run_strategy(
                         name=field_name,
                         kind=inspect.Parameter.KEYWORD_ONLY,
                         default=OptionInfo(
-                            default=field_meta.default,
+                            default=field_info.default,
                             param_decls=list(set([f"--{field_name}", f"--{field_name.replace('_', '-')}"])),
-                            help=f"{field_meta.field_info.description}",
+                            help=f"{field_info.description}",
                             rich_help_panel="Strategy Settings",
                         ),
-                        annotation=__process_type(field_meta.type_),
+                        annotation=__process_type(field_info.annotation),
                     )
-                    for field_name, field_meta in strategy_type.get_settings_type().__fields__.items()
+                    for field_name, field_info in strategy_type.get_settings_type().model_fields.items()
                 ]
             )