Add advanced exclude and include support for dict, json and copy

HISTORY.rst

@@ -5,6 +5,7 @@ History
 
 v0.31 (unreleased)
 ..................
+* add advanced exclude support for ``dict``, ``json`` and ``copy``, #648 by @MrMrRobat
 * nested classes which inherit and change ``__init__`` are now correctly processed while still allowing ``self`` as a
   parameter, #644 by @lnaden and @dgasmith
 

pydantic/main.py

@@ -37,6 +37,7 @@
     AnyType,
     ForwardRef,
     GetterDict,
+    ValueItems,
     change_exception,
     is_classvar,
     resolve_annotations,
@@ -58,7 +59,9 @@
     SetStr = Set[str]
     ListStr = List[str]
     Model = TypeVar('Model', bound='BaseModel')
-
+    IntStr = Union[int, str]
+    SetIntStr = Set[IntStr]
+    DictIntStrAny = Dict[IntStr, Any]
 
 try:
     import cython  # type: ignore
@@ -304,21 +307,30 @@ def __setstate__(self, state: 'DictAny') -> None:
         object.__setattr__(self, '__fields_set__', state['__fields_set__'])
 
     def dict(
-        self, *, include: 'SetStr' = None, exclude: 'SetStr' = None, by_alias: bool = False, skip_defaults: bool = False
+        self,
+        *,
+        include: Union['SetIntStr', 'DictIntStrAny'] = None,
+        exclude: Union['SetIntStr', 'DictIntStrAny'] = None,
+        by_alias: bool = False,
+        skip_defaults: bool = False,
     ) -> 'DictStrAny':
         """
         Generate a dictionary representation of the model, optionally specifying which fields to include or exclude.
         """
         get_key = self._get_key_factory(by_alias)
         get_key = partial(get_key, self.fields)
 
-        return_keys = self._calculate_keys(include=include, exclude=exclude, skip_defaults=skip_defaults)
-        if return_keys is None:
-            return {get_key(k): v for k, v in self._iter(by_alias=by_alias, skip_defaults=skip_defaults)}
-        else:
-            return {
-                get_key(k): v for k, v in self._iter(by_alias=by_alias, skip_defaults=skip_defaults) if k in return_keys
-            }
+        allowed_keys = self._calculate_keys(include=include, exclude=exclude, skip_defaults=skip_defaults)
+        return {
+            get_key(k): v
+            for k, v in self._iter(
+                by_alias=by_alias,
+                allowed_keys=allowed_keys,
+                include=include,
+                exclude=exclude,
+                skip_defaults=skip_defaults,
+            )
+        }
 
     def _get_key_factory(self, by_alias: bool) -> Callable[..., str]:
         if by_alias:
@@ -329,8 +341,8 @@ def _get_key_factory(self, by_alias: bool) -> Callable[..., str]:
     def json(
         self,
         *,
-        include: 'SetStr' = None,
-        exclude: 'SetStr' = None,
+        include: Union['SetIntStr', 'DictIntStrAny'] = None,
+        exclude: Union['SetIntStr', 'DictIntStrAny'] = None,
         by_alias: bool = False,
         skip_defaults: bool = False,
         encoder: Optional[Callable[[Any], Any]] = None,
@@ -417,8 +429,8 @@ def construct(cls: Type['Model'], values: 'DictAny', fields_set: 'SetStr') -> 'M
     def copy(
         self: 'Model',
         *,
-        include: 'SetStr' = None,
-        exclude: 'SetStr' = None,
+        include: Union['SetIntStr', 'DictIntStrAny'] = None,
+        exclude: Union['SetIntStr', 'DictIntStrAny'] = None,
         update: 'DictStrAny' = None,
         deep: bool = False,
     ) -> 'Model':
@@ -436,11 +448,23 @@ def copy(
             # skip constructing values if no arguments are passed
             v = self.__values__
         else:
-            return_keys = self._calculate_keys(include=include, exclude=exclude, skip_defaults=False)
-            if return_keys:
-                v = {**{k: v for k, v in self.__values__.items() if k in return_keys}, **(update or {})}
-            else:
+            allowed_keys = self._calculate_keys(include=include, exclude=exclude, skip_defaults=False, update=update)
+            if allowed_keys is None:
                 v = {**self.__values__, **(update or {})}
+            else:
+                v = {
+                    **dict(
+                        self._iter(
+                            to_dict=False,
+                            by_alias=False,
+                            include=include,
+                            exclude=exclude,
+                            skip_defaults=False,
+                            allowed_keys=allowed_keys,
+                        )
+                    ),
+                    **(update or {}),
+                }
 
         if deep:
             v = deepcopy(v)
@@ -487,17 +511,56 @@ def _decompose_class(cls: Type['Model'], obj: Any) -> GetterDict:
         return GetterDict(obj)
 
     @classmethod
-    def _get_value(cls, v: Any, by_alias: bool, skip_defaults: bool) -> Any:
+    @no_type_check
+    def _get_value(
+        cls,
+        v: Any,
+        to_dict: bool,
+        by_alias: bool,
+        include: Optional[Union['SetIntStr', 'DictIntStrAny']],
+        exclude: Optional[Union['SetIntStr', 'DictIntStrAny']],
+        skip_defaults: bool,
+    ) -> Any:
+
         if isinstance(v, BaseModel):
-            return v.dict(by_alias=by_alias, skip_defaults=skip_defaults)
-        elif isinstance(v, list):
-            return [cls._get_value(v_, by_alias=by_alias, skip_defaults=skip_defaults) for v_ in v]
-        elif isinstance(v, dict):
-            return {k_: cls._get_value(v_, by_alias=by_alias, skip_defaults=skip_defaults) for k_, v_ in v.items()}
-        elif isinstance(v, set):
-            return {cls._get_value(v_, by_alias=by_alias, skip_defaults=skip_defaults) for v_ in v}
-        elif isinstance(v, tuple):
-            return tuple(cls._get_value(v_, by_alias=by_alias, skip_defaults=skip_defaults) for v_ in v)
+            if to_dict:
+                return v.dict(by_alias=by_alias, skip_defaults=skip_defaults, include=include, exclude=exclude)
+            else:
+                return v.copy(include=include, exclude=exclude)
+
+        value_exclude = ValueItems(v, exclude) if exclude else None
+        value_include = ValueItems(v, include) if include else None
+
+        if isinstance(v, dict):
+            return {
+                k_: cls._get_value(
+                    v_,
+                    to_dict=to_dict,
+                    by_alias=by_alias,
+                    skip_defaults=skip_defaults,
+                    include=value_include and value_include.for_element(k_),
+                    exclude=value_exclude and value_exclude.for_element(k_),
+                )
+                for k_, v_ in v.items()
+                if (not value_exclude or not value_exclude.is_excluded(k_))
+                and (not value_include or value_include.is_included(k_))
+            }
+
+        elif isinstance(v, (list, set, tuple)):
+            return type(v)(
+                cls._get_value(
+                    v_,
+                    to_dict=to_dict,
+                    by_alias=by_alias,
+                    skip_defaults=skip_defaults,
+                    include=value_include and value_include.for_element(i),
+                    exclude=value_exclude and value_exclude.for_element(i),
+                )
+                for i, v_ in enumerate(v)
+                if (not value_exclude or not value_exclude.is_excluded(i))
+                and (not value_include or value_include.is_included(i))
+            )
+
         else:
             return v
 
@@ -517,14 +580,37 @@ def __iter__(self) -> 'AnyGenerator':
         """
         yield from self._iter()
 
-    def _iter(self, by_alias: bool = False, skip_defaults: bool = False) -> 'TupleGenerator':
+    def _iter(
+        self,
+        to_dict: bool = True,
+        by_alias: bool = False,
+        allowed_keys: Optional['SetStr'] = None,
+        include: Union['SetIntStr', 'DictIntStrAny'] = None,
+        exclude: Union['SetIntStr', 'DictIntStrAny'] = None,
+        skip_defaults: bool = False,
+    ) -> 'TupleGenerator':
+
+        value_exclude = ValueItems(self, exclude) if exclude else None
+        value_include = ValueItems(self, include) if include else None
+
         for k, v in self.__values__.items():
-            yield k, self._get_value(v, by_alias=by_alias, skip_defaults=skip_defaults)
+            if allowed_keys is None or k in allowed_keys:
+                yield k, self._get_value(
+                    v,
+                    to_dict=to_dict,
+                    by_alias=by_alias,
+                    include=value_include and value_include.for_element(k),
+                    exclude=value_exclude and value_exclude.for_element(k),
+                    skip_defaults=skip_defaults,
+                )
 
     def _calculate_keys(
-        self, include: 'SetStr' = None, exclude: Optional['SetStr'] = None, skip_defaults: bool = False
+        self,
+        include: Optional[Union['SetIntStr', 'DictIntStrAny']] = None,
+        exclude: Optional[Union['SetIntStr', 'DictIntStrAny']] = None,
+        skip_defaults: bool = False,
+        update: Optional['DictStrAny'] = None,
     ) -> Optional['SetStr']:
-
         if include is None and exclude is None and skip_defaults is False:
             return None
 
@@ -533,11 +619,18 @@ def _calculate_keys(
         else:
             keys = set(self.__values__.keys())
 
-        if include:
-            keys &= include
+        if include and isinstance(include, dict):
+            keys &= set(include.keys())
+        elif include:
+            keys &= include  # type: ignore
+
+        if update:
+            keys -= set(update.keys())
 
-        if exclude:
-            keys -= exclude
+        if exclude and isinstance(exclude, dict):
+            keys -= {k for k, v in exclude.items() if v is ...}
+        elif exclude:
+            keys -= exclude  # type: ignore
 
         return keys
 

pydantic/utils.py

@@ -20,6 +20,7 @@
     Type,
     Union,
     _eval_type,
+    no_type_check,
 )
 
 import pydantic
@@ -48,6 +49,7 @@
 if TYPE_CHECKING:  # pragma: no cover
     from .main import BaseModel  # noqa: F401
     from .main import Field  # noqa: F401
+    from .main import SetIntStr, DictIntStrAny, IntStr  # noqa: F401
     from . import errors  # noqa: F401
 
 if sys.version_info < (3, 7):
@@ -350,3 +352,82 @@ def keys(self) -> Set[Any]:
         We don't want to get any other attributes of obj if the model didn't explicitly ask for them
         """
         return set()
+
+
+class ValueItems:
+    """
+    Class for more convenient calculation of excluded or included fields on values.
+    """
+
+    __slots__ = ('_items', '_type')
+
+    def __init__(self, value: Any, items: Union['SetIntStr', 'DictIntStrAny']) -> None:
+        if TYPE_CHECKING:
+            self._items: Union['SetIntStr', 'DictIntStrAny']
+            self._type: Type[Union[set, dict]]  # type: ignore
+
+        # For further type checks speed-up
+        if isinstance(items, dict):
+            self._type = dict
+        elif isinstance(items, set):
+            self._type = set
+        else:
+            raise ValueError(f'Unexpected type of exclude value {type(items)}')
+
+        if isinstance(value, (list, tuple)):
+            items = self._normalize_indexes(items, len(value))
+
+        self._items = items
+
+    @no_type_check
+    def is_excluded(self, item: Any) -> bool:
+        """
+        Check if item is fully excluded
+        (value considered excluded if self._type is set and item contained in self._items
+         or self._type is dict and self._items.get(item) is ...
+
+        :param item: key or index of a value
+        """
+        if self._type is set:
+            return item in self._items
+        return self._items.get(item) is ...
+
+    @no_type_check
+    def is_included(self, item: Any) -> bool:
+        """
+        Check if value is contained in self._items
+
+        :param item: key or index of value
+        """
+        return item in self._items
+
+    @no_type_check
+    def for_element(self, e: 'IntStr') -> Optional[Union['SetIntStr', 'DictIntStrAny']]:
+        """
+        :param e: key or index of element on value
+        :return: raw values for elemet if self._items is dict and contain needed element
+        """
+
+        if self._type is dict:
+            item = self._items.get(e)
+            return item if item is not ... else None
+        return None
+
+    @no_type_check
+    def _normalize_indexes(
+        self, items: Union['SetIntStr', 'DictIntStrAny'], v_length: int
+    ) -> Union['SetIntStr', 'DictIntStrAny']:
+        """
+        :param items: dict or set of indexes which will be normalized
+        :param v_length: length of sequence indexes of which will be
+
+        >>> self._normalize_indexes({0, -2, -1}, 4)
+        {0, 2, 3}
+        """
+        if self._type is set:
+            return {v_length + i if i < 0 else i for i in items}
+        else:
+            return {v_length + i if i < 0 else i: v for i, v in items.items()}
+
+    def __str__(self) -> str:
+        return f'{self.__class__.__name__}: {self._type.__name__}({self._items})'