sub-models working

pydantic · samuelcolvin · Jun 28, 2018 · Jun 3, 2018 · Jun 3, 2018 · Jun 4, 2018
commit 31b237e63dfe2fd83957f248a3de641be404c74f
diff --git a/pydantic/fields.py b/pydantic/fields.py
@@ -33,7 +33,7 @@ class Validator(NamedTuple):
 
 
 class Schema:
-    __slots__ = 'alias', 'title', 'description', 'choice_names', 'extra',
+    __slots__ = 'alias', 'title', 'choice_names', 'extra',
 
     def __init__(self, alias=None, title=None, choice_names=None, **extra):
         self.alias = alias
@@ -85,10 +85,9 @@ def infer(cls, *, name, value, annotation, class_validators, config):
         schema_from_config = config.get_field_schema(name)
        if isinstance(value, tuple) and len(value) == 2 and isinstance(value[1], Schema):
             value, schema = value
-            schema.alias = schema.alias or schema_from_config.get('alias')
         else:
             schema = Schema(**schema_from_config)
-        schema.title = schema.title or name.title()
+        schema.alias = schema.alias or schema_from_config.get('alias')
         required = value == Required
         return cls(
             name=name,
@@ -129,20 +128,21 @@ def prepare(self):
         self._populate_sub_fields()
         self._populate_validators()
 
-    def schema(self):
-        s = {
-            'type': display_as_type(self.type_),
-            'required': self.required,
-            'title': self._schema.title,
-        }
-        if not self.required:
+    def schema(self, by_alias=True):
+        s = self.type_.schema(by_alias) if hasattr(self.type_, 'schema') else {}
+        s.update(
+            type=s.get('type') or display_as_type(self.type_),
+            title=self._schema.title or s.get('title') or self.alias.title(),
+            required=self.required,
+        )
+
+        if not self.required and self.default is not None:
             s['default'] = self.default
-        if isinstance(self.type_, Enum):
+        if issubclass(self.type_, Enum):
             if self._schema.choice_names:
                 s['choices'] = [(v.value, self._schema.choice_names[v.value]) for v in self.type_.__members__.values()]
             else:
                 s['choices'] = [(v.value, k) for k, v in self.type_.__members__.items()]
-        # TODO gt, lg, sub-model properties
         s.update(self._schema.extra)
         return s
 

diff --git a/pydantic/main.py b/pydantic/main.py
@@ -11,11 +11,12 @@
 from .fields import Field, Validator
 from .parse import Protocol, load_file, load_str_bytes
 from .types import StrBytes
-from .utils import truncate
+from .utils import clean_docstring, truncate
 from .validators import dict_validator
 
 
 class BaseConfig:
+    title = None
     anystr_strip_whitespace = False
     min_anystr_length = 0
     max_anystr_length = 2 ** 16
@@ -132,6 +133,7 @@ def __new__(mcs, name, bases, namespace):
             '__config__': config,
             '__fields__': fields,
             '__validators__': vg.validators,
+            '_schema_cache': {},
             **{n: v for n, v in namespace.items() if n not in fields}
         }
         return super().__new__(mcs, name, bases, new_namespace)
@@ -250,6 +252,25 @@ def copy(self, *, include: Set[str]=None, exclude: Set[str]=None, update: Dict[s
     def fields(self):
         return self.__fields__
 
+    @classmethod
+    def schema(cls, by_alias=True):
+        cached = cls._schema_cache.get(by_alias)
+        if cached is not None:
+            return cached
+        if by_alias:
+            props = {f.alias: f.schema(by_alias) for f in cls.__fields__.values()}
+        else:
+            props = {k: f.schema(by_alias) for k, f in cls.__fields__.items()}
+        s = {
+            'type': 'object',
+            'title': cls.__config__.title or cls.__name__,
+            'properties': props,
+        }
+        if cls.__doc__:
+            s['description'] = clean_docstring(cls.__doc__)
+        cls._schema_cache[by_alias] = s
+        return s
+
     @classmethod
     def get_validators(cls):
         yield dict_validator

diff --git a/pydantic/utils.py b/pydantic/utils.py
@@ -1,6 +1,8 @@
 import re
 from contextlib import contextmanager
+from enum import Enum
 from importlib import import_module
+    from textwrap import dedent
 from typing import Tuple, _TypingBase
 
 from . import errors
@@ -118,6 +120,14 @@ def display_as_type(v):
     if not isinstance(v, _TypingBase) and not isinstance(v, type):
         v = type(v)
 
+    if isinstance(v, type) and issubclass(v, Enum):
+        if issubclass(v, int):
+            return 'int'
+        elif issubclass(v, str):
+            return 'str'
+        else:
+            return 'enum'
+
     try:
         return v.__name__
     except AttributeError:
@@ -131,3 +141,7 @@ def change_exception(raise_exc, *except_types):
         yield
     except except_types as e:
         raise raise_exc from e
+
+
+def clean_docstring(d):
+    return dedent(d).strip(' \r\n\t')
diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -0,0 +1,138 @@
+from enum import Enum, IntEnum
+
+from pydantic import BaseModel
+
+
+def test_key():
+    class ApplePie(BaseModel):
+        """
+        This is a test.
+        """
+        a: float
+        b: int = 10
+
+    s = {
+        'type': 'object',
+        'title': 'ApplePie',
+        'description': 'This is a test.',
+        'properties': {
+            'a': {
+                'type': 'float',
+                'required': True,
+                'title': 'A',
+            },
+            'b': {
+                'type': 'int',
+                'required': False,
+                'title': 'B',
+                'default': 10,
+            },
+        },
+    }
+    assert True not in ApplePie._schema_cache
+    assert False not in ApplePie._schema_cache
+    assert ApplePie.schema() == s
+    assert True in ApplePie._schema_cache
+    assert False not in ApplePie._schema_cache
+    assert ApplePie.schema() == s
+
+
+def test_by_alias():
+    class ApplePie(BaseModel):
+        a: float
+        b: int = 10
+
+        class Config:
+            title = 'Apple Pie'
+            fields = {'a': 'Snap', 'b': 'Crackle'}
+
+    s = {
+        'type': 'object',
+        'title': 'Apple Pie',
+        'properties': {
+            'Snap': {
+                'type': 'float',
+                'required': True,
+                'title': 'Snap',
+            },
+            'Crackle': {
+                'type': 'int',
+                'required': False,
+                'title': 'Crackle',
+                'default': 10,
+            },
+        },
+    }
+    assert ApplePie.schema() == s
+    assert ApplePie.schema() == s
+    assert list(ApplePie.schema(by_alias=True)['properties'].keys()) == ['Snap', 'Crackle']
+    assert list(ApplePie.schema(by_alias=False)['properties'].keys()) == ['a', 'b']
+
+
+def test_sub_model():
+    class Foo(BaseModel):
+        """hello"""
+        b: float
+
+    class Bar(BaseModel):
+        a: int
+        b: Foo = None
+
+    assert Bar.schema() == {
+        'type': 'object',
+        'title': 'Bar',
+        'properties': {
+            'a': {
+                'type': 'int',
+                'title': 'A',
+                'required': True,
+            },
+            'b': {
+                'type': 'object',
+                'title': 'Foo',
+                'properties': {
+                    'b': {
+                        'type': 'float',
+                        'title': 'B',
+                        'required': True,
+                    },
+                },
+                'description': 'hello',
+                'required': False,
+            },
+        },
+    }
+
+
+def test_choices():
+    FooEnum = Enum('FooEnum', {'foo': 'f', 'bar': 'b'})
+    BarEnum = IntEnum('BarEnum', {'foo': 1, 'bar': 2})
+
+    class Model(BaseModel):
+        foo: FooEnum
+        bar: BarEnum
+
+    assert Model.schema() == {
+        'type': 'object',
+        'title': 'Model',
+        'properties': {
+            'foo': {
+                'type': 'enum',
+                'title': 'Foo',
+                'required': True,
+                'choices': [
+                    ('f', 'foo'),
+                    ('b', 'bar'),
+                ],
+                },
+            'bar': {
+                'type': 'int',
+                'title': 'Bar',
+                'required': True,
+                'choices': [
+                    (1, 'foo'),
+                    (2, 'bar'),
+                ],
+            },
+        },
+    }