From 19c0a391e02e075871da0747036e4b80e5864dd0 Mon Sep 17 00:00:00 2001
From: Mike Oakley <18501112+mike-oakley@users.noreply.github.com>
Date: Fri, 13 May 2022 17:41:38 +0100
Subject: [PATCH] Add datatype enum.

---
 openapi_schema_pydantic/v3/v3_0_3/datatype.py | 12 +++++++
 openapi_schema_pydantic/v3/v3_0_3/schema.py   |  4 ++-
 openapi_schema_pydantic/v3/v3_1_0/datatype.py | 13 +++++++
 openapi_schema_pydantic/v3/v3_1_0/schema.py   |  3 +-
 tests/v3_0_3/test_datatype.py                 | 34 ++++++++++++++++++
 tests/v3_1_0/test_datatype.py                 | 35 +++++++++++++++++++
 6 files changed, 99 insertions(+), 2 deletions(-)
 create mode 100644 openapi_schema_pydantic/v3/v3_0_3/datatype.py
 create mode 100644 openapi_schema_pydantic/v3/v3_1_0/datatype.py
 create mode 100644 tests/v3_0_3/test_datatype.py
 create mode 100644 tests/v3_1_0/test_datatype.py

diff --git a/openapi_schema_pydantic/v3/v3_0_3/datatype.py b/openapi_schema_pydantic/v3/v3_0_3/datatype.py
new file mode 100644
index 0000000..abf3523
--- /dev/null
+++ b/openapi_schema_pydantic/v3/v3_0_3/datatype.py
@@ -0,0 +1,12 @@
+import enum
+
+
+class DataType(str, enum.Enum):
+    """Data type of an object."""
+
+    STRING = "string"
+    NUMBER = "number"
+    INTEGER = "integer"
+    BOOLEAN = "boolean"
+    ARRAY = "array"
+    OBJECT = "object"
diff --git a/openapi_schema_pydantic/v3/v3_0_3/schema.py b/openapi_schema_pydantic/v3/v3_0_3/schema.py
index d29bd36..b7428d3 100644
--- a/openapi_schema_pydantic/v3/v3_0_3/schema.py
+++ b/openapi_schema_pydantic/v3/v3_0_3/schema.py
@@ -1,6 +1,8 @@
 from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Extra, Field
+
+from .datatype import DataType
 from .discriminator import Discriminator
 from .external_documentation import ExternalDocumentation
 from .reference import Reference
@@ -206,7 +208,7 @@ class Schema(BaseModel):
     but their definitions were adjusted to the OpenAPI Specification.
     """
 
-    type: Optional[str] = None
+    type: Optional[DataType] = None
     """
     **From OpenAPI spec:
     Value MUST be a string. Multiple types via an array are not supported.**
diff --git a/openapi_schema_pydantic/v3/v3_1_0/datatype.py b/openapi_schema_pydantic/v3/v3_1_0/datatype.py
new file mode 100644
index 0000000..6477eee
--- /dev/null
+++ b/openapi_schema_pydantic/v3/v3_1_0/datatype.py
@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class DataType(str, Enum):
+    """Data type of an object."""
+
+    NULL = "null"
+    STRING = "string"
+    NUMBER = "number"
+    INTEGER = "integer"
+    BOOLEAN = "boolean"
+    ARRAY = "array"
+    OBJECT = "object"
diff --git a/openapi_schema_pydantic/v3/v3_1_0/schema.py b/openapi_schema_pydantic/v3/v3_1_0/schema.py
index 1739f0e..61e45c2 100644
--- a/openapi_schema_pydantic/v3/v3_1_0/schema.py
+++ b/openapi_schema_pydantic/v3/v3_1_0/schema.py
@@ -1,6 +1,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Extra, Field
+from .datatype import DataType
 from .discriminator import Discriminator
 from .external_documentation import ExternalDocumentation
 from .reference import Reference
@@ -358,7 +359,7 @@ class Schema(BaseModel):
     and follow the same specifications:
     """
 
-    type: Optional[Union[str, List[str]]] = None
+    type: Optional[Union[DataType, List[DataType]]] = None
     """
     The value of this keyword MUST be either a string or an array.  If it
     is an array, elements of the array MUST be strings and MUST be
diff --git a/tests/v3_0_3/test_datatype.py b/tests/v3_0_3/test_datatype.py
new file mode 100644
index 0000000..3d61749
--- /dev/null
+++ b/tests/v3_0_3/test_datatype.py
@@ -0,0 +1,34 @@
+import pytest
+from pydantic import ValidationError
+
+from openapi_schema_pydantic.v3.v3_0_3 import Schema
+
+
+@pytest.mark.parametrize(
+    "datatype",
+    (
+        "string",
+        "number",
+        "integer",
+        "boolean",
+        "array",
+        "object",
+    ),
+)
+def test_good_types_parse_and_equate(datatype: str):
+    assert Schema(type=datatype).type == datatype
+
+
+def test_bad_types_raise_validation_errors():
+    with pytest.raises(ValidationError):
+        Schema(type="invalid")
+
+    with pytest.raises(ValidationError):
+        Schema(anyOf=[{"type": "invalid"}])
+
+    with pytest.raises(ValidationError):
+        Schema(
+            properties={
+                "a": Schema(type="invalid"),
+            },
+        )
diff --git a/tests/v3_1_0/test_datatype.py b/tests/v3_1_0/test_datatype.py
new file mode 100644
index 0000000..b1ffc31
--- /dev/null
+++ b/tests/v3_1_0/test_datatype.py
@@ -0,0 +1,35 @@
+import pytest
+from pydantic import ValidationError
+
+from openapi_schema_pydantic.v3.v3_1_0 import Schema
+
+
+@pytest.mark.parametrize(
+    "datatype",
+    (
+        "string",
+        "number",
+        "integer",
+        "boolean",
+        "array",
+        "object",
+        "null",
+    ),
+)
+def test_good_types_parse_and_equate(datatype: str):
+    assert Schema(type=datatype).type == datatype
+
+
+def test_bad_types_raise_validation_errors():
+    with pytest.raises(ValidationError):
+        Schema(type="invalid")
+
+    with pytest.raises(ValidationError):
+        Schema(anyOf=[{"type": "invalid"}])
+
+    with pytest.raises(ValidationError):
+        Schema(
+            properties={
+                "a": Schema(type="invalid"),
+            },
+        )