From 6b5226c74a45de82a7eb5462953ba8c8a0fc9574 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Voron?= <fvoron@gmail.com>
Date: Thu, 5 May 2022 08:32:02 +0200
Subject: [PATCH] Add IntegerIdMixin and export related classes

---
 fastapi_users/__init__.py |  6 ++++++
 fastapi_users/manager.py  | 10 ++++++++++
 tests/test_manager.py     | 18 ++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/fastapi_users/__init__.py b/fastapi_users/__init__.py
index 6d177186..13ecb84f 100644
--- a/fastapi_users/__init__.py
+++ b/fastapi_users/__init__.py
@@ -7,6 +7,9 @@ from fastapi_users.fastapi_users import FastAPIUsers  # noqa: F401
 from fastapi_users.manager import (  # noqa: F401
     BaseUserManager,
     InvalidPasswordException,
+    InvalidID,
+    UUIDIDMixin,
+    IntegerIDMixin,
 )
 
 __all__ = [
@@ -14,4 +17,7 @@ __all__ = [
     "FastAPIUsers",
     "BaseUserManager",
     "InvalidPasswordException",
+    "InvalidID",
+    "UUIDIDMixin",
+    "IntegerIDMixin",
 ]
diff --git a/fastapi_users/manager.py b/fastapi_users/manager.py
index 116db23a..a4d1c1cd 100644
--- a/fastapi_users/manager.py
+++ b/fastapi_users/manager.py
@@ -612,4 +612,14 @@ class UUIDIDMixin:
             raise InvalidID() from e
 
 
+class IntegerIDMixin:
+    def parse_id(self, value: Any) -> int:
+        if isinstance(value, float):
+            raise InvalidID()
+        try:
+            return int(value)
+        except ValueError as e:
+            raise InvalidID() from e
+
+
 UserManagerDependency = DependencyCallable[BaseUserManager[models.UP, models.ID]]
diff --git a/tests/test_manager.py b/tests/test_manager.py
index 6f423aab..634668cf 100644
--- a/tests/test_manager.py
+++ b/tests/test_manager.py
@@ -7,6 +7,8 @@ from pytest_mock import MockerFixture
 
 from fastapi_users.jwt import decode_jwt, generate_jwt
 from fastapi_users.manager import (
+    IntegerIDMixin,
+    InvalidID,
     InvalidPasswordException,
     InvalidResetPasswordToken,
     InvalidVerifyToken,
@@ -591,3 +593,19 @@ class TestAuthenticate:
         assert user is not None
         assert user.email == "king.arthur@camelot.bt"
         assert update_spy.called is True
+
+
+def test_integer_id_mixin():
+    integer_id_mixin = IntegerIDMixin()
+
+    assert integer_id_mixin.parse_id("123") == 123
+    assert integer_id_mixin.parse_id(123) == 123
+
+    with pytest.raises(InvalidID):
+        integer_id_mixin.parse_id("123.42")
+
+    with pytest.raises(InvalidID):
+        integer_id_mixin.parse_id(123.42)
+
+    with pytest.raises(InvalidID):
+        integer_id_mixin.parse_id("abc")