From 2f39bc4f4f6f459e3c63e72d87b65e01b617ec6b Mon Sep 17 00:00:00 2001 From: Rami Abdelrazzaq Date: Sat, 7 Mar 2026 23:08:52 -0600 Subject: [PATCH] Improve typing for Database.__getitem__ --- sqlite_utils/db.py | 4 ++-- tests/test_typing.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) create mode 100644 tests/test_typing.py diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index aacdc893..dc1ea50b 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -445,7 +445,7 @@ def tracer( finally: self._tracer = prev_tracer - def __getitem__(self, table_name: str) -> Union["Table", "View"]: + def __getitem__(self, table_name: str) -> "Table": """ ``db[table_name]`` returns a :class:`.Table` object for the table with the specified name. If the table does not exist yet it will be created the first time data is inserted into it. @@ -453,7 +453,7 @@ def __getitem__(self, table_name: str) -> Union["Table", "View"]: :param table_name: The name of the table """ if table_name in self.view_names(): - return self.view(table_name) + return cast("Table", self.view(table_name)) return self.table(table_name) def __repr__(self) -> str: diff --git a/tests/test_typing.py b/tests/test_typing.py new file mode 100644 index 00000000..cf32a365 --- /dev/null +++ b/tests/test_typing.py @@ -0,0 +1,13 @@ +import typing + +from sqlite_utils.db import Database, Table, View + + +def test_database_getitem_return_type_hint_is_table(): + assert typing.get_type_hints(Database.__getitem__)["return"] is Table + + +def test_database_getitem_still_returns_view_for_views(fresh_db): + fresh_db["items"].insert({"name": "one"}) + fresh_db.create_view("items_view", "select * from items") + assert isinstance(fresh_db["items_view"], View)