Add model subscription lookup, Model[<pkval>] that will return the ob…

…ject or raise KeyError (#386)

CHANGELOG.rst

@@ -14,6 +14,7 @@ Changelog
 - fix: ``sqlite://:memory:`` in Windows thrown ``OSError: [WinError 123]``
 - Support ``bulk_create()`` insertion of records with overridden primary key when the primary key is DB-generated
 - Add ``queryset.exists()`` and ``Model.exists()``.
+- Add model subscription lookup, ``Model[<pkval>]`` that will return the object or raise ``KeyError``
 
 0.16.10
 -------

tests/test_model_methods.py

@@ -206,6 +206,18 @@ async def test_update_from_dict(self):
         with self.assertRaises(ValueError):
             evt2.update_from_dict({"alias": "foo"})
 
+    async def test_index_access(self):
+        obj = await self.cls[self.mdl.pk]
+        self.assertEqual(obj, self.mdl)
+
+    async def test_index_badval(self):
+        with self.assertRaises(KeyError):
+            await self.cls[100000]
+
+    async def test_index_badtype(self):
+        with self.assertRaises(KeyError):
+            await self.cls["asdf"]
+
 
 class TestModelMethodsNoID(TestModelMethods):
     async def setUp(self):

tortoise/models.py

@@ -22,6 +22,7 @@
 from tortoise.backends.base.client import BaseDBAsyncClient
 from tortoise.exceptions import (
     ConfigurationError,
+    DoesNotExist,
     IncompleteInstanceError,
     IntegrityError,
     OperationalError,
@@ -602,6 +603,9 @@ def __search_for_field_attributes(base: Type, attrs: dict) -> None:
         meta.finalise_fields()
         return new_class
 
+    def __getitem__(cls: Type[MODEL], key: Any) -> QuerySetSingle[MODEL]:  # type: ignore
+        return cls._getbypk(key)  # type: ignore
+
 
 class Model(metaclass=ModelMeta):
     """
@@ -728,6 +732,13 @@ def _set_pk_val(self, value: Any) -> None:
     Can be used as a field name when doing filtering e.g. ``.filter(pk=...)`` etc...
     """
 
+    @classmethod
+    async def _getbypk(cls: Type[MODEL], key: Any) -> MODEL:
+        try:
+            return await cls.get(pk=key)
+        except (DoesNotExist, ValueError):
+            raise KeyError(f"{cls._meta.full_name} has no object {repr(key)}")
+
     def update_from_dict(self, data: dict) -> MODEL:
         """
         Updates the current model with the provided dict.