From d91ef7ecf921c372a9122ce83023b7c1f476844c Mon Sep 17 00:00:00 2001 From: Paul Yang Date: Mon, 10 Feb 2025 12:43:04 -0800 Subject: [PATCH] Update handling for Pydantic objects in `mf_pformat` (#1666) This PR updates the handling for Pydantic objects in `mf_pformat` to check if the `dict()` methods exists instead of the `isinstance(BaseModel)` check. The `BaseModel` is actually from the `dsi_shim` and in some cases with different Pydantic versions, it's possible that the check does not pass. Since it's not foolproof, the formatter falls back to `pprint` if there is an error. --- .../mf_logging/pretty_print.py | 42 ++++++++++--------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/metricflow-semantics/metricflow_semantics/mf_logging/pretty_print.py b/metricflow-semantics/metricflow_semantics/mf_logging/pretty_print.py index 20801b3bb..3e6c5be27 100644 --- a/metricflow-semantics/metricflow_semantics/mf_logging/pretty_print.py +++ b/metricflow-semantics/metricflow_semantics/mf_logging/pretty_print.py @@ -2,12 +2,9 @@ import logging import pprint -from collections.abc import Mapping from dataclasses import fields, is_dataclass from enum import Enum -from typing import Any, List, Optional, Sized, Tuple, Union - -from dsi_pydantic_shim import BaseModel +from typing import Any, List, Mapping, Optional, Sized, Tuple, Union from metricflow_semantics.mf_logging.formatting import indent from metricflow_semantics.mf_logging.pretty_formattable import MetricFlowPrettyFormattable @@ -35,11 +32,6 @@ def __init__( self._include_none_object_fields = include_none_object_fields self._include_empty_object_fields = include_empty_object_fields - @staticmethod - def _is_pydantic_base_model(obj: Any): # type:ignore - # Checking the attribute as the BaseModel check fails for newer version of Pydantic. - return isinstance(obj, BaseModel) or hasattr(obj, "__config__") - def _handle_sequence_obj( self, list_like_obj: Union[list, tuple, set, frozenset], remaining_line_width: Optional[int] ) -> str: @@ -353,16 +345,28 @@ def _handle_any_obj(self, obj: Any, remaining_line_width: Optional[int]) -> str: remaining_line_width=remaining_line_width, ) - if MetricFlowPrettyFormatter._is_pydantic_base_model(obj): - mapping = {key: getattr(obj, key) for key in obj.dict().keys()} - return self._handle_mapping_like_obj( - mapping, - left_enclose_str=type(obj).__name__ + "(", - key_value_seperator="=", - right_enclose_str=")", - is_dataclass_like_object=True, - remaining_line_width=remaining_line_width, - ) + # For Pydantic-like objects with a `dict`-like method that returns field keys / values. + # In Pydantic v1, it's `.dict()`, in v2 it's `.model_dump()`. + # Going with this approach for now to check for a Pydantic model as using `isinstance()` requires more + # consideration when dealing with Pydantic v1 and Pydantic v2 objects. + pydantic_dict_method = getattr(obj, "model_dump", None) or getattr(obj, "dict", None) + if pydantic_dict_method is not None and callable(pydantic_dict_method): + try: + # Calling `dict` on a Pydantic model does not recursively convert nested fields into dictionaries, + # which is what we want. `.model_dump()` does the recursive conversion. + mapping = dict(obj) + return self._handle_mapping_like_obj( + mapping, + left_enclose_str=type(obj).__name__ + "(", + key_value_seperator="=", + right_enclose_str=")", + is_dataclass_like_object=True, + remaining_line_width=remaining_line_width, + ) + except Exception: + # Fall back to built-in pretty-print in case the dict method can't be called. e.g. requires arguments. + # Consider logging a warning. + pass # Any other object that's not handled. return pprint.pformat(obj, width=self._max_line_width, sort_dicts=False)