From e77bc00d6e52a4ea13fb467c163ca67468c7d02c Mon Sep 17 00:00:00 2001
From: Andrey Golovizin <ag+services@sologoc.com>
Date: Thu, 21 Feb 2019 23:23:50 +0100
Subject: [PATCH] Fix type hints of `parse_obj` and similar methods (#405)

---
 HISTORY.rst                |  4 ++++
 pydantic/main.py           | 25 ++++++++++++++++---------
 tests/mypy_test_success.py | 27 +++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index ca8adc2..ba2f4d0 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -3,6 +3,10 @@
 History
 -------
 
+v0.20.1 (unreleased)
+....................
+* fix type hints of ``parse_obj`` and similar methods, #405 by @erosennin
+
 v0.20.0 (2019-02-18)
 ....................
 * fix tests for python 3.8, #396 by @samuelcolvin
diff --git a/pydantic/main.py b/pydantic/main.py
index d13b088..82734e8 100644
--- a/pydantic/main.py
+++ b/pydantic/main.py
@@ -19,6 +19,7 @@ from typing import (
     Set,
     Tuple,
     Type,
+    TypeVar,
     Union,
     cast,
     no_type_check,
@@ -53,6 +54,7 @@ if TYPE_CHECKING:  # pragma: no cover
     DictAny = Dict[Any, Any]
     SetStr = Set[str]
     ListStr = List[str]
+    Model = TypeVar('Model', bound='BaseModel')
 
 
 class Extra(str, Enum):
@@ -307,7 +309,7 @@ class BaseModel(metaclass=MetaModel):
         )
 
     @classmethod
-    def parse_obj(cls, obj: 'DictAny') -> 'BaseModel':
+    def parse_obj(cls: Type['Model'], obj: 'DictAny') -> 'Model':
         if not isinstance(obj, dict):
             exc = TypeError(f'{cls.__name__} expected dict not {type(obj).__name__}')
             raise ValidationError([ErrorWrapper(exc, loc='__obj__')])
@@ -315,14 +317,14 @@ class BaseModel(metaclass=MetaModel):
 
     @classmethod
     def parse_raw(
-        cls,
+        cls: Type['Model'],
         b: StrBytes,
         *,
         content_type: str = None,
         encoding: str = 'utf8',
         proto: Protocol = None,
         allow_pickle: bool = False,
-    ) -> 'BaseModel':
+    ) -> 'Model':
         try:
             obj = load_str_bytes(
                 b, proto=proto, content_type=content_type, encoding=encoding, allow_pickle=allow_pickle
@@ -333,19 +335,19 @@ class BaseModel(metaclass=MetaModel):
 
     @classmethod
     def parse_file(
-        cls,
+        cls: Type['Model'],
         path: Union[str, Path],
         *,
         content_type: str = None,
         encoding: str = 'utf8',
         proto: Protocol = None,
         allow_pickle: bool = False,
-    ) -> 'BaseModel':
+    ) -> 'Model':
         obj = load_file(path, proto=proto, content_type=content_type, encoding=encoding, allow_pickle=allow_pickle)
         return cls.parse_obj(obj)
 
     @classmethod
-    def construct(cls, values: 'DictAny', fields_set: 'SetStr') -> 'BaseModel':
+    def construct(cls: Type['Model'], values: 'DictAny', fields_set: 'SetStr') -> 'Model':
         """
         Creates a new model and set __values__ without any validation, thus values should already be trusted.
         Chances are you don't want to use this method directly.
@@ -356,8 +358,13 @@ class BaseModel(metaclass=MetaModel):
         return m
 
     def copy(
-        self, *, include: 'SetStr' = None, exclude: 'SetStr' = None, update: 'DictStrAny' = None, deep: bool = False
-    ) -> 'BaseModel':
+        self: 'Model',
+        *,
+        include: 'SetStr' = None,
+        exclude: 'SetStr' = None,
+        update: 'DictStrAny' = None,
+        deep: bool = False,
+    ) -> 'Model':
         """
         Duplicate a model, optionally choose which fields to include, exclude and change.
 
@@ -407,7 +414,7 @@ class BaseModel(metaclass=MetaModel):
         yield cls.validate
 
     @classmethod
-    def validate(cls, value: Union['DictStrAny', 'BaseModel']) -> 'BaseModel':
+    def validate(cls: Type['Model'], value: Union['DictStrAny', 'Model']) -> 'Model':
         if isinstance(value, dict):
             return cls(**value)
         elif isinstance(value, BaseModel):
diff --git a/tests/mypy_test_success.py b/tests/mypy_test_success.py
index 10a61dc..714f1f8 100644
--- a/tests/mypy_test_success.py
+++ b/tests/mypy_test_success.py
@@ -3,6 +3,7 @@ Test pydantic's compliance with mypy.
 
 Do a little skipping about with types to demonstrate its usage.
 """
+import json
 from datetime import datetime
 from typing import List, Optional
 
@@ -47,6 +48,32 @@ assert m.signup_ts == datetime(2017, 6, 7), m.signup_ts
 assert day_of_week(m.signup_ts) == 3
 
 
+data = {'age': 10, 'first_name': 'Alena', 'last_name': 'Sousova', 'list_of_ints': [410]}
+m_from_obj = Model.parse_obj(data)
+
+assert isinstance(m_from_obj, Model)
+assert m_from_obj.age == 10
+assert m_from_obj.first_name == data['first_name']
+assert m_from_obj.last_name == data['last_name']
+assert m_from_obj.list_of_ints == data['list_of_ints']
+
+m_from_raw = Model.parse_raw(json.dumps(data))
+
+assert isinstance(m_from_raw, Model)
+assert m_from_raw.age == m_from_obj.age
+assert m_from_raw.first_name == m_from_obj.first_name
+assert m_from_raw.last_name == m_from_obj.last_name
+assert m_from_raw.list_of_ints == m_from_obj.list_of_ints
+
+m_copy = m_from_obj.copy()
+
+assert isinstance(m_from_raw, Model)
+assert m_copy.age == m_from_obj.age
+assert m_copy.first_name == m_from_obj.first_name
+assert m_copy.last_name == m_from_obj.last_name
+assert m_copy.list_of_ints == m_from_obj.list_of_ints
+
+
 @dataclass
 class AddProject:
     name: str