diff --git a/synapse/types/__init__.py b/synapse/types/__init__.py index 6621145f25..90d7c0ed72 100644 --- a/synapse/types/__init__.py +++ b/synapse/types/__init__.py @@ -27,6 +27,7 @@ from enum import Enum from typing import ( TYPE_CHECKING, AbstractSet, + Annotated, Any, ClassVar, Final, @@ -42,10 +43,11 @@ from typing import ( overload, ) +import annotated_types import attr import pydantic_core.core_schema from immutabledict import immutabledict -from pydantic import GetCoreSchemaHandler +from pydantic import GetCoreSchemaHandler, StrictInt from pydantic_core import CoreSchema from signedjson.key import decode_verify_key_bytes from signedjson.types import VerifyKey @@ -174,6 +176,13 @@ For a Sentinel for internal (non-API-facing) use, instead consider """ +NonNegativeStrictInt = Annotated[StrictInt, annotated_types.Ge(0)] +"""A strict integer that must be greater than or equal to zero. + +Should be preferred in place of Pydantic's own (lax) NonNegativeInt. +""" + + # Note that this seems to require inheriting *directly* from Interface in order # for mypy-zope to realize it is an interface. class ISynapseThreadlessReactor( diff --git a/tests/test_types.py b/tests/test_types.py index fb8735d8a4..43fd96d6f5 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -31,6 +31,7 @@ from synapse.types import ( AbsentType, AbstractMultiWriterStreamToken, MultiWriterStreamToken, + NonNegativeStrictInt, RoomAlias, RoomStreamToken, UserID, @@ -278,3 +279,74 @@ class AbsentTestCase(unittest.TestCase): self.assertIs(copy.copy(Absent), Absent) self.assertIs(a.absent, b.absent) + + +class NonNegativeStrictIntTestCase(unittest.TestCase): + """ + Tests for the `NonNegativeStrictInt` utility. + """ + + def test_pydantic_jsonschema(self) -> None: + """ + Tests that `NonNegativeStrictInt` produces sensible JSONSchema. + """ + + class MyModel(BaseModel): + limit: NonNegativeStrictInt = 100 + + self.assertEqual( + MyModel.model_json_schema(), + { + "properties": { + "limit": { + "default": 100, + "minimum": 0, + "title": "Limit", + "type": "integer", + } + }, + "title": "MyModel", + "type": "object", + }, + f"JSONSchema actually is:\n{MyModel.model_json_schema()!r}", + ) + + def test_pydantic_reject(self) -> None: + """ + Tests that `NonNegativeStrictInt` rejects negative numbers + and non-ints. + """ + + class MyModel(BaseModel): + limit: NonNegativeStrictInt = 100 + + with self.assertRaises(ValidationError): + MyModel.model_validate({"limit": -1}) + + with self.assertRaises(ValidationError): + MyModel.model_validate_json('{"limit": -1}') + + # StrictInt, so don't accept floats... + with self.assertRaises(ValidationError): + MyModel.model_validate({"limit": 1.5}) + + with self.assertRaises(ValidationError): + MyModel.model_validate_json('{"limit": 1.5}') + + # ...and don't accept stringy ints either. + with self.assertRaises(ValidationError): + MyModel.model_validate({"limit": "42"}) + + with self.assertRaises(ValidationError): + MyModel.model_validate_json('{"limit": "42"}') + + def test_pydantic_accept(self) -> None: + """ + Tests that `Absent` accepts the absence of a value when used in Pydantic models. + """ + + class MyModel(BaseModel): + limit: NonNegativeStrictInt = 100 + + self.assertEqual(MyModel.model_validate_json('{"limit": 0}'), MyModel(limit=0)) + self.assertEqual(MyModel.model_validate({"limit": 42}), MyModel(limit=42))