feat: add SecretStr and SecretBytes. (#452)

* feat: add SecretStr and SecretBytes. * chore: update HISTORY.rst * fix: file permissions were incorrect. * feat: lint, format, fix comments. * feat: changed inner type of SecretBytes in the schema to string as there is no bytes type in json. * feat: remove format from secret str and secret bytes. * feat: fix schema mapping.
pydantic · Apr 4, 2019 · 4a8faca · 4a8faca
1 parent 694abaf
commit 4a8faca
Show file tree

Hide file tree

Showing 9 changed files with 197 additions and 1 deletion.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -13,6 +13,7 @@ v0.23 (unreleased)
 * fix ``ForwardRef`` collection bug, #450 by @tigerwings
 * Support specialized ``ClassVars``, #455 by @tyrylu
 * fix JSON serialization for ``ipaddress`` types, #333 by @pilosus
+* add SecretStr and SecretBytes types, #452 by @atheuz
 
 v0.22 (2019-03-29)
 ....................

diff --git a/docs/examples/ex_secret_types.py b/docs/examples/ex_secret_types.py
@@ -0,0 +1,29 @@
+from typing import List
+
+from pydantic import BaseModel, SecretStr, SecretBytes, ValidationError
+
+class SimpleModel(BaseModel):
+    password: SecretStr
+    password_bytes: SecretBytes
+
+print(SimpleModel(password='IAmSensitive', password_bytes=b'IAmSensitiveBytes'))
+# > SimpleModel password=SecretStr('**********') password_bytes=SecretBytes(b'**********')
+
+sm = SimpleModel(password='IAmSensitive', password_bytes=b'IAmSensitiveBytes')
+print(sm.password.get_secret_value())
+# > IAmSensitive
+print(sm.password_bytes.get_secret_value())
+# > b'IAmSensitiveBytes'
+
+
+try:
+    SimpleModel(password=[1,2,3], password_bytes=[1,2,3])
+except ValidationError as e:
+    print(e)
+"""
+2 validation error
+password
+  str type expected (type=type_error.str)
+password_bytes
+  byte type expected (type=type_error.bytes)
+"""
diff --git a/docs/examples/exotic.py b/docs/examples/exotic.py
@@ -6,7 +6,7 @@
 
 from pydantic import (DSN, UUID1, UUID3, UUID4, UUID5, BaseModel, DirectoryPath, EmailStr, FilePath, NameEmail,
                       NegativeFloat, NegativeInt, PositiveFloat, PositiveInt, PyObject, UrlStr, conbytes, condecimal,
-                      confloat, conint, constr, IPvAnyAddress, IPvAnyInterface, IPvAnyNetwork)
+                      confloat, conint, constr, IPvAnyAddress, IPvAnyInterface, IPvAnyNetwork, SecretStr, SecretBytes)
 
 
 class Model(BaseModel):
@@ -39,6 +39,9 @@ class Model(BaseModel):
 
     url: UrlStr = None
 
+    password: SecretStr = None
+    password_bytes: SecretBytes = None
+
     db_name = 'foobar'
     db_user = 'postgres'
     db_password: str = None
@@ -89,6 +92,8 @@ class Model(BaseModel):
     email_address='Samuel Colvin <s@muelcolvin.com >',
     email_and_name='Samuel Colvin <s@muelcolvin.com >',
     url='http://example.com',
+    password='password',
+    password_bytes=b'password2'
     decimal=Decimal('42.24'),
     decimal_positive=Decimal('21.12'),
     decimal_negative=Decimal('-21.12'),
@@ -133,6 +138,8 @@ class Model(BaseModel):
     'email_address': 's@muelcolvin.com',
     'email_and_name': <NameEmail("Samuel Colvin <s@muelcolvin.com>")>,
     'url': 'http://example.com',
+    'password': SecretStr('**********'),
+    'password_bytes': SecretStr(b'**********'),
     ...
     'dsn': 'postgres://postgres@localhost:5432/foobar',
     'decimal': Decimal('42.24'),

diff --git a/docs/index.rst b/docs/index.rst
@@ -458,6 +458,16 @@ Fields can also be of type ``Callable``:
     callable, no validation of arguments, their types or the return
     type is performed.
 
+Secret Types
+............
+
+You can use the ``SecretStr`` and the ``SecretBytes`` data types for storing sensitive information
+that you do not want to be visible in logging or tracebacks.
+
+.. literalinclude:: examples/ex_secret_types.py
+
+(This script is complete, it should run "as is")
+
 Json Type
 .........
 

diff --git a/docs/schema_mapping.py b/docs/schema_mapping.py
@@ -107,6 +107,20 @@
         'JSON Schema Validation',
         'All the literal values in the enum are included in the definition.'
     ],
+    [
+        'SecretStr',
+        'string',
+        '{"writeOnly": true}',
+        'JSON Schema Validation',
+        ''
+    ],
+    [
+        'SecretBytes',
+        'string',
+        '{"writeOnly": true}',
+        'JSON Schema Validation',
+        ''
+    ],
     [
         'EmailStr',
         'string',

diff --git a/pydantic/schema.py b/pydantic/schema.py
@@ -24,6 +24,8 @@
     FilePath,
     Json,
     NameEmail,
+    SecretBytes,
+    SecretStr,
     UrlStr,
     condecimal,
     confloat,
@@ -587,7 +589,9 @@ def field_singleton_sub_fields_schema(
     (EmailStr, {'type': 'string', 'format': 'email'}),
     (UrlStr, {'type': 'string', 'format': 'uri'}),
     (DSN, {'type': 'string', 'format': 'dsn'}),
+    (SecretStr, {'type': 'string', 'writeOnly': True}),
     (str, {'type': 'string'}),
+    (SecretBytes, {'type': 'string', 'writeOnly': True}),
     (bytes, {'type': 'string', 'format': 'binary'}),
     (bool, {'type': 'boolean'}),
     (int, {'type': 'integer'}),

diff --git a/pydantic/types.py b/pydantic/types.py
@@ -74,6 +74,8 @@
     'IPvAnyAddress',
     'IPvAnyInterface',
     'IPvAnyNetwork',
+    'SecretStr',
+    'SecretBytes',
 ]
 
 NoneStr = Optional[str]
@@ -569,3 +571,49 @@ def validate(cls, value: NetworkType) -> Union[IPv4Network, IPv6Network]:
 
         with change_exception(errors.IPvAnyNetworkError, ValueError):
             return IPv6Network(value)
+
+
+class SecretStr:
+    @classmethod
+    def __get_validators__(cls) -> 'CallableGenerator':
+        yield str_validator
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, value: str) -> 'SecretStr':
+        return cls(value)
+
+    def __init__(self, value: str):
+        self._secret_value = value
+
+    def __repr__(self) -> str:
+        return "SecretStr('**********')" if self._secret_value else "SecretStr('')"
+
+    def __str__(self) -> str:
+        return repr(self)
+
+    def get_secret_value(self) -> str:
+        return self._secret_value
+
+
+class SecretBytes:
+    @classmethod
+    def __get_validators__(cls) -> 'CallableGenerator':
+        yield bytes_validator
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, value: bytes) -> 'SecretBytes':
+        return cls(value)
+
+    def __init__(self, value: bytes):
+        self._secret_value = value
+
+    def __repr__(self) -> str:
+        return "SecretBytes(b'**********')" if self._secret_value else "SecretBytes(b'')"
+
+    def __str__(self) -> str:
+        return repr(self)
+
+    def get_secret_value(self) -> bytes:
+        return self._secret_value
diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -36,6 +36,8 @@
     PositiveFloat,
     PositiveInt,
     PyObject,
+    SecretBytes,
+    SecretStr,
     StrBytes,
     StrictStr,
     UrlStr,
@@ -517,6 +519,21 @@ class Model(BaseModel):
     assert Model.schema() == base_schema
 
 
+@pytest.mark.parametrize('field_type,inner_type', [(SecretBytes, 'string'), (SecretStr, 'string')])
+def test_secret_types(field_type, inner_type):
+    class Model(BaseModel):
+        a: field_type
+
+    base_schema = {
+        'title': 'Model',
+        'type': 'object',
+        'properties': {'a': {'title': 'A', 'type': inner_type, 'writeOnly': True}},
+        'required': ['a'],
+    }
+
+    assert Model.schema() == base_schema
+
+
 @pytest.mark.parametrize(
     'field_type,expected_schema',
     [

diff --git a/tests/test_types.py b/tests/test_types.py
@@ -28,6 +28,8 @@
     PositiveFloat,
     PositiveInt,
     PyObject,
+    SecretBytes,
+    SecretStr,
     StrictStr,
     ValidationError,
     conbytes,
@@ -1411,3 +1413,67 @@ class Foobar(BaseModel):
     assert exc_info.value.errors() == [
         {'loc': ('pattern',), 'msg': 'Invalid regular expression', 'type': 'value_error.regex_pattern'}
     ]
+
+
+def test_secretstr():
+    class Foobar(BaseModel):
+        password: SecretStr
+        empty_password: SecretStr
+
+    # Initialize the model.
+    f = Foobar(password='1234', empty_password='')
+
+    # Assert correct types.
+    assert f.password.__class__.__name__ == 'SecretStr'
+    assert f.empty_password.__class__.__name__ == 'SecretStr'
+
+    # Assert str and repr are correct.
+    assert str(f.password) == "SecretStr('**********')"
+    assert str(f.empty_password) == "SecretStr('')"
+    assert repr(f.password) == "SecretStr('**********')"
+    assert repr(f.empty_password) == "SecretStr('')"
+
+    # Assert retrieval of secret value is correct
+    assert f.password.get_secret_value() == '1234'
+    assert f.empty_password.get_secret_value() == ''
+
+
+def test_secretstr_error():
+    class Foobar(BaseModel):
+        password: SecretStr
+
+    with pytest.raises(ValidationError) as exc_info:
+        Foobar(password=[6, 23, 'abc'])
+    assert exc_info.value.errors() == [{'loc': ('password',), 'msg': 'str type expected', 'type': 'type_error.str'}]
+
+
+def test_secretbytes():
+    class Foobar(BaseModel):
+        password: SecretBytes
+        empty_password: SecretBytes
+
+    # Initialize the model.
+    f = Foobar(password=b'wearebytes', empty_password=b'')
+
+    # Assert correct types.
+    assert f.password.__class__.__name__ == 'SecretBytes'
+    assert f.empty_password.__class__.__name__ == 'SecretBytes'
+
+    # Assert str and repr are correct.
+    assert str(f.password) == "SecretBytes(b'**********')"
+    assert str(f.empty_password) == "SecretBytes(b'')"
+    assert repr(f.password) == "SecretBytes(b'**********')"
+    assert repr(f.empty_password) == "SecretBytes(b'')"
+
+    # Assert retrieval of secret value is correct
+    assert f.password.get_secret_value() == b'wearebytes'
+    assert f.empty_password.get_secret_value() == b''
+
+
+def test_secretbytes_error():
+    class Foobar(BaseModel):
+        password: SecretBytes
+
+    with pytest.raises(ValidationError) as exc_info:
+        Foobar(password=[6, 23, 'abc'])
+    assert exc_info.value.errors() == [{'loc': ('password',), 'msg': 'byte type expected', 'type': 'type_error.bytes'}]