fix code, use pyotp

pydantic_extra_types/__init__.py

@@ -1,5 +1,5 @@
 __version__ = '0.0.1'
 
-from pydantic_extra_types.types import OTP, OTP_ALPHABET, Color, PaymentCardBrand, PaymentCardNumber
+from pydantic_extra_types.types import Color, OTPToken, PaymentCardBrand, PaymentCardNumber
 
-__all__ = 'Color', 'PaymentCardNumber', 'PaymentCardBrand', 'OTP', 'OTP_ALPHABET'
+__all__ = 'Color', 'PaymentCardNumber', 'PaymentCardBrand', 'OTPToken'

pydantic_extra_types/types/__init__.py

@@ -1,5 +1,5 @@
 from pydantic_extra_types.types.color import Color
-from pydantic_extra_types.types.otp import OTP, OTP_ALPHABET
+from pydantic_extra_types.types.otp import OTPToken
 from pydantic_extra_types.types.payment import PaymentCardBrand, PaymentCardNumber
 
-__all__ = 'Color', 'PaymentCardNumber', 'PaymentCardBrand', 'OTP', 'OTP_ALPHABET'
+__all__ = 'Color', 'PaymentCardNumber', 'PaymentCardBrand', 'OTPToken'

pydantic_extra_types/types/otp.py

@@ -1,43 +1,29 @@
-from typing import Any, ClassVar
+from typing import Any
 
-from pydantic_core import PydanticCustomError, core_schema
+import pyotp
+from pydantic_core import PydanticCustomError
 
-# Default OTP Number of Digits
-OTP_ALPHABET = '0123456789'
 
+class OTPToken(str):
+    """A one-time password token.
 
-class OTP(str):
-    """
-    One-time password
-    """
-
-    strip_whitespace: ClassVar[bool] = True
-    min_length: ClassVar[int] = 6
-    max_length: ClassVar[int] = 6
-
-    def __init__(self, otp: str):
-        self.validate_digits(otp)
-        self.validate_length(otp)
+    This is a custom type that can be used to validate a one-time password token
+    against a secret key. The secret key is passed in the context argument of
+    the model_validate method.
 
-    @classmethod
-    def __get_pydantic_core_schema__(cls, **_kwargs: Any) -> core_schema.FunctionSchema:
-        return core_schema.function_after_schema(
-            core_schema.str_schema(
-                min_length=cls.min_length, max_length=cls.max_length, strip_whitespace=cls.strip_whitespace
-            ),
-            cls.validate,
-        )
+    The type also has a custom JSON encoder that returns the current one-time
+    password token for the secret key.
+    """
 
-    @classmethod
-    def validate(cls, __input_value: str, **_kwargs: Any) -> 'OTP':
-        return cls(__input_value)
+    @staticmethod
+    def model_validate(value: Any, *, context: Any) -> Any:
+        if not pyotp.TOTP(context['otp_secret']).verify(value):
+            raise PydanticCustomError('Invalid one-time password', value)
+        return value
 
     @classmethod
-    def validate_digits(cls, otp: str) -> None:
-        if not otp.isdigit():
-            raise PydanticCustomError('otp_digits', 'OTP is not all digits')
+    def get_validators(cls) -> Any:
+        yield cls.model_validate
 
-    @classmethod
-    def validate_length(cls, otp: str) -> None:
-        if len(otp) != cls.max_length:
-            raise PydanticCustomError('otp_length', 'OTP must be {length} digits', {'length': cls.max_length})
+    class Config:
+        json_encoders = {pyotp.TOTP: lambda v: v.now()}

pyproject.toml

@@ -39,6 +39,7 @@ classifiers = [
 requires-python = '>=3.7'
 dependencies = [
     'pydantic@git+https://github.com/pydantic/pydantic.git@main',
+    'pyotp',
 ]
 dynamic = ['version']
 

requirements/pyproject.txt

@@ -10,6 +10,8 @@ pydantic @ git+https://github.com/pydantic/pydantic.git@main
     # via pydantic-extra-types (pyproject.toml)
 pydantic-core==0.9.0
     # via pydantic
+pyotp==2.8.0
+    # via pydantic-extra-types (pyproject.toml)
 typing-extensions==4.4.0
     # via
     #   pydantic

tests/test_types_otp.py

@@ -1,30 +1,35 @@
+import pyotp
 import pytest
 from pydantic_core import PydanticCustomError
 
-from pydantic_extra_types import OTP
+from pydantic_extra_types import OTPToken
 
 
-def test_otp():
-    # Test successful validation
-    valid_otp = OTP.validate('123456')
-    assert isinstance(valid_otp, OTP)
+def test_model_validate():
+    secret = 'JBSWY3DPEHPK3PXP'
+    totp = pyotp.TOTP(secret)
+    value = totp.now()
+    context = {'otp_secret': secret}
+    # Test a valid OTP token
+    result = OTPToken.model_validate(value, context=context)
+    assert result == value
 
-    # Test validation error for non-digit characters
-    with pytest.raises(PydanticCustomError) as excinfo:
-        OTP.validate('1A2345')
-    assert str(excinfo.value) == 'OTP is not all digits'
+    # Test an invalid OTP token
+    with pytest.raises(PydanticCustomError, match='Invalid one-time password'):
+        OTPToken.model_validate('Invalid one-time password', context=context)
 
-    # Test validation error for length less than 6
-    with pytest.raises(PydanticCustomError) as excinfo:
-        OTP.validate('12345')
-    assert str(excinfo.value) == 'OTP must be 6 digits'
 
-    # Test validation error for length greater than 6
-    with pytest.raises(PydanticCustomError) as excinfo:
-        OTP.validate('1234567')
-    assert str(excinfo.value) == 'OTP must be 6 digits'
+def test_json_encoder():
+    secret = 'JBSWY3DPEHPK3PXP'
+    totp = pyotp.TOTP(secret)
+    value = totp.now()
 
-    # Test validation error for empty string
-    with pytest.raises(PydanticCustomError) as excinfo:
-        OTP.validate('')
-    assert str(excinfo.value) == 'OTP is not all digits'
+    # Test the custom JSON encoder
+    json_encoded_value = OTPToken.Config.json_encoders[pyotp.TOTP](totp)
+    assert json_encoded_value == value
+
+
+def test_get_validators():
+    # Test the get_validators method
+    validators = list(OTPToken.get_validators())
+    assert validators == [OTPToken.model_validate]