From 18f4111ce436f2a0e812ead3ad08e996974c0b93 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Thu, 4 Aug 2022 11:33:32 -0400 Subject: [PATCH] Use the correct type annotations for PrefectBaseModel.subclass() (#6205) * Use the correct type annotations for PrefectBaseModel.subclass() --- src/prefect/orion/utilities/schemas.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/prefect/orion/utilities/schemas.py b/src/prefect/orion/utilities/schemas.py index 1fd92cc4a0bf..1e217a01ce26 100644 --- a/src/prefect/orion/utilities/schemas.py +++ b/src/prefect/orion/utilities/schemas.py @@ -6,7 +6,7 @@ import json import os from functools import partial -from typing import Any, Dict, List, Set, TypeVar +from typing import Any, Dict, List, Set, Type, TypeVar from uuid import UUID, uuid4 import pendulum @@ -31,12 +31,15 @@ def validate(cls, v): raise ValueError("Unrecognized datetime.") +B = TypeVar("B", bound=BaseModel) + + def pydantic_subclass( - base: BaseModel, + base: Type[B], name: str = None, include_fields: List[str] = None, exclude_fields: List[str] = None, -) -> BaseModel: +) -> Type[B]: """Creates a subclass of a Pydantic model that excludes certain fields. Pydantic models use the __fields__ attribute of their parent class to determine inherited fields, so to create a subclass without fields, we @@ -137,11 +140,11 @@ class Config: @classmethod def subclass( - cls, + cls: Type[B], name: str = None, include_fields: List[str] = None, exclude_fields: List[str] = None, - ) -> BaseModel: + ) -> Type[B]: """Creates a subclass of this model containing only the specified fields. See `pydantic_subclass()`.