From 7b7ebecffe253e835ecd0051cc5ed365ae2b638e Mon Sep 17 00:00:00 2001
From: Amin Alaee <mohammadamin.alaee@gmail.com>
Date: Sun, 10 Oct 2021 12:06:55 +0200
Subject: [PATCH] add IPAddress field (#114)

---
 docs/fields.md         | 11 +++++++
 tests/test_fields.py   | 28 ++++++++++++++++++
 tests/test_schemas.py  | 32 +++++++++++++++++----
 typesystem/__init__.py |  2 ++
 typesystem/fields.py   |  6 ++++
 typesystem/formats.py  | 65 +++++++++++++++++++++++++++++++++++++-----
 6 files changed, 132 insertions(+), 12 deletions(-)

diff --git a/docs/fields.md b/docs/fields.md
index f534d53..311a0fe 100644
--- a/docs/fields.md
+++ b/docs/fields.md
@@ -83,6 +83,17 @@ Represented in HTML forms as a `<textarea>`.
 Similar to `String` and takes the same arguments.
 Represented in HTML forms as a `<input type="password">`.
 
+### Email
+
+Similar to `String` and takes the same arguments.
+Represented in HTML forms as a `<input type="email">`.
+
+### IPAddress
+
+Validates IPv4 and IPv6 addresses.
+
+Returns `ipaddress.IPv4Address` or `ipaddress.IPv6Address` based on input.
+
 ## Boolean data types
 
 ### Boolean
diff --git a/tests/test_fields.py b/tests/test_fields.py
index 91a1d47..e57d5fa 100644
--- a/tests/test_fields.py
+++ b/tests/test_fields.py
@@ -1,5 +1,6 @@
 import datetime
 import decimal
+import ipaddress
 import re
 import uuid
 
@@ -16,6 +17,7 @@
     Email,
     Float,
     Integer,
+    IPAddress,
     Number,
     Object,
     Password,
@@ -850,3 +852,29 @@ def test_password():
     validator = Password()
     value, _ = validator.validate_or_error("secret")
     assert value == "secret"
+
+
+def test_ipaddress():
+    validator = IPAddress()
+    value, error = validator.validate_or_error("192.168.1.1")
+    assert value == ipaddress.ip_address("192.168.1.1")
+
+    validator = IPAddress()
+    value, error = validator.validate_or_error("192.168.1.af")
+    assert error == ValidationError(text="Must be a valid IP format.", code="format")
+
+    validator = IPAddress()
+    value, error = validator.validate_or_error("192.168.1.256")
+    assert error == ValidationError(text="Must be a real IP.", code="invalid")
+
+    validator = IPAddress()
+    value, error = validator.validate_or_error(
+        "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
+    )
+    assert value == ipaddress.ip_address("2001:0db8:85a3:0000:0000:8a2e:0370:7334")
+
+    validator = IPAddress()
+    value, error = validator.validate_or_error(
+        "2001:0dz8:85a3:0000:0000:8a2e:0370:7334"
+    )
+    assert error == ValidationError(text="Must be a valid IP format.", code="format")
diff --git a/tests/test_schemas.py b/tests/test_schemas.py
index e3c31ea..37cd432 100644
--- a/tests/test_schemas.py
+++ b/tests/test_schemas.py
@@ -1,4 +1,6 @@
 import datetime
+import ipaddress
+import uuid
 
 import typesystem
 import typesystem.formats
@@ -265,16 +267,22 @@ def test_schema_decimal_serialization():
 def test_schema_uuid_serialization():
     user = typesystem.Schema(
         fields={
-            "id": typesystem.String(format="uuid"),
+            "id": typesystem.UUID(),
+            "parent_id": typesystem.UUID(),
             "username": typesystem.String(),
         }
     )
 
-    item = {"id": "b769df4a-18ec-480f-89ef-8ea961a82269", "username": "tom"}
+    item = {
+        "id": uuid.UUID("b769df4a-18ec-480f-89ef-8ea961a82269"),
+        "username": "tom",
+        "parent_id": None,
+    }
     data = user.serialize(item)
 
     assert data["id"] == "b769df4a-18ec-480f-89ef-8ea961a82269"
     assert data["username"] == "tom"
+    assert data["parent_id"] is None
 
 
 def test_schema_reference_serialization():
@@ -503,8 +511,22 @@ def test_definitions_to_json_schema():
 
 
 def test_schema_email_serialization():
-    user = typesystem.Schema(fields={"email": typesystem.Email()})
+    user = typesystem.Schema(
+        fields={"from": typesystem.Email(), "to": typesystem.Email()}
+    )
 
-    item = {"email": "team@encode.io"}
+    item = {"from": "team@encode.io", "to": None}
     data = user.serialize(item)
-    assert data == {"email": "team@encode.io"}
+    assert data["from"] == "team@encode.io"
+    assert data["to"] is None
+
+
+def test_schema_ipaddress_serialization():
+    schema = typesystem.Schema(
+        fields={"src": typesystem.IPAddress(), "dst": typesystem.IPAddress()}
+    )
+
+    item = {"src": ipaddress.ip_address("127.0.0.1"), "dst": None}
+    data = schema.serialize(item)
+    assert data["src"] == "127.0.0.1"
+    assert data["dst"] is None
diff --git a/typesystem/__init__.py b/typesystem/__init__.py
index 1f72729..b7bc206 100644
--- a/typesystem/__init__.py
+++ b/typesystem/__init__.py
@@ -12,6 +12,7 @@
     Field,
     Float,
     Integer,
+    IPAddress,
     Number,
     Object,
     Password,
@@ -38,6 +39,7 @@
     "Decimal",
     "Email",
     "Integer",
+    "IPAddress",
     "Jinja2Forms",
     "Field",
     "Float",
diff --git a/typesystem/fields.py b/typesystem/fields.py
index b78255a..ca0fb28 100644
--- a/typesystem/fields.py
+++ b/typesystem/fields.py
@@ -15,6 +15,7 @@
     "datetime": formats.DateTimeFormat(),
     "uuid": formats.UUIDFormat(),
     "email": formats.EmailFormat(),
+    "ipaddress": formats.IPAddressFormat(),
 }
 
 
@@ -778,3 +779,8 @@ def __init__(self, **kwargs: typing.Any) -> None:
 class Password(String):
     def __init__(self, **kwargs: typing.Any) -> None:
         super().__init__(format="password", **kwargs)
+
+
+class IPAddress(String):
+    def __init__(self, **kwargs: typing.Any) -> None:
+        super().__init__(format="ipaddress", **kwargs)
diff --git a/typesystem/formats.py b/typesystem/formats.py
index ed8c5cc..e6f94c5 100644
--- a/typesystem/formats.py
+++ b/typesystem/formats.py
@@ -1,4 +1,5 @@
 import datetime
+import ipaddress
 import re
 import typing
 import uuid
@@ -30,6 +31,13 @@
     re.IGNORECASE,
 )
 
+IPV4_REGEX = re.compile(
+    r"(?:0|25[0-5]|2[0-4]\d|1\d?\d?|[1-9]\d?)"
+    r"(?:\.(?:0|25[0-5]|2[0-4]\d|1\d?\d?|[1-9]\d?)){3}"
+)
+
+IPV6_REGEX = re.compile(r"(?:[a-f0-9]{1,4}:){7}[a-f0-9]{1,4}")
+
 
 class BaseFormat:
     errors: typing.Dict[str, str] = {}
@@ -44,7 +52,7 @@ def is_native_type(self, value: typing.Any) -> bool:
     def validate(self, value: typing.Any) -> typing.Union[typing.Any, ValidationError]:
         raise NotImplementedError()  # pragma: no cover
 
-    def serialize(self, obj: typing.Any) -> typing.Union[str, None]:
+    def serialize(self, obj: typing.Any) -> typing.Optional[str]:
         raise NotImplementedError()  # pragma: no cover
 
 
@@ -68,7 +76,7 @@ def validate(self, value: typing.Any) -> datetime.date:
         except ValueError:
             raise self.validation_error("invalid")
 
-    def serialize(self, obj: typing.Any) -> typing.Union[str, None]:
+    def serialize(self, obj: typing.Optional[datetime.date]) -> typing.Optional[str]:
         if obj is None:
             return None
 
@@ -101,7 +109,7 @@ def validate(self, value: typing.Any) -> datetime.time:
         except ValueError:
             raise self.validation_error("invalid")
 
-    def serialize(self, obj: typing.Any) -> typing.Union[str, None]:
+    def serialize(self, obj: typing.Optional[datetime.time]) -> typing.Optional[str]:
         if obj is None:
             return None
 
@@ -147,7 +155,9 @@ def validate(self, value: typing.Any) -> datetime.datetime:
         except ValueError:
             raise self.validation_error("invalid")
 
-    def serialize(self, obj: typing.Any) -> typing.Union[str, None]:
+    def serialize(
+        self, obj: typing.Optional[datetime.datetime]
+    ) -> typing.Optional[str]:
         if obj is None:
             return None
 
@@ -174,7 +184,12 @@ def validate(self, value: typing.Any) -> uuid.UUID:
 
         return uuid.UUID(value)
 
-    def serialize(self, obj: typing.Any) -> str:
+    def serialize(self, obj: typing.Optional[uuid.UUID]) -> typing.Optional[str]:
+        if obj is None:
+            return None
+
+        assert isinstance(obj, uuid.UUID)
+
         return str(obj)
 
 
@@ -184,12 +199,48 @@ class EmailFormat(BaseFormat):
     def is_native_type(self, value: typing.Any) -> bool:
         return False
 
-    def validate(self, value: typing.Any) -> uuid.UUID:
+    def validate(self, value: str) -> str:
         match = EMAIL_REGEX.match(value)
         if not match:
             raise self.validation_error("format")
 
         return value
 
-    def serialize(self, obj: typing.Any) -> str:
+    def serialize(self, obj: typing.Optional[str]) -> typing.Optional[str]:
+        if obj is None:
+            return None
+
+        return obj
+
+
+class IPAddressFormat(BaseFormat):
+    errors = {
+        "format": "Must be a valid IP format.",
+        "invalid": "Must be a real IP.",
+    }
+
+    def is_native_type(self, value: typing.Any) -> bool:
+        return isinstance(value, (ipaddress.IPv4Address, ipaddress.IPv6Address))
+
+    def validate(
+        self, value: typing.Any
+    ) -> typing.Union[ipaddress.IPv4Address, ipaddress.IPv6Address]:
+        match_ipv4 = IPV4_REGEX.match(value)
+        match_ipv6 = IPV6_REGEX.match(value)
+        if not match_ipv4 and not match_ipv6:
+            raise self.validation_error("format")
+
+        try:
+            return ipaddress.ip_address(value)
+        except ValueError:
+            raise self.validation_error("invalid")
+
+    def serialize(
+        self, obj: typing.Union[ipaddress.IPv4Address, ipaddress.IPv6Address, None]
+    ) -> typing.Optional[str]:
+        if obj is None:
+            return None
+
+        assert isinstance(obj, (ipaddress.IPv4Address, ipaddress.IPv6Address))
+
         return str(obj)