db.col_name and db['col_name'] aliases for get_collection

datastax · Feb 29, 2024 · 97ad75a · 97ad75a
1 parent 17e7c8a
commit 97ad75a
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 6 deletions.
diff --git a/astrapy/idiomatic/database.py b/astrapy/idiomatic/database.py
@@ -113,9 +113,11 @@ def __init__(
             caller_version=caller_version,
         )
 
-    @property
-    def namespace(self) -> str:
-        return self._astra_db.namespace
+    def __getattr__(self, collection_name: str) -> Collection:
+        return self.get_collection(name=collection_name)
+
+    def __getitem__(self, collection_name: str) -> Collection:
+        return self.get_collection(name=collection_name)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}[_astra_db={self._astra_db}"]'
@@ -126,6 +128,10 @@ def __eq__(self, other: Any) -> bool:
         else:
             return False
 
+    @property
+    def namespace(self) -> str:
+        return self._astra_db.namespace
+
     def copy(self) -> Database:
         return Database(
             **self._constructor_params,
@@ -324,9 +330,11 @@ def __init__(
             caller_version=caller_version,
         )
 
-    @property
-    def namespace(self) -> str:
-        return self._astra_db.namespace
+    async def __getattr__(self, collection_name: str) -> AsyncCollection:
+        return await self.get_collection(name=collection_name)
+
+    async def __getitem__(self, collection_name: str) -> AsyncCollection:
+        return await self.get_collection(name=collection_name)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}[_astra_db={self._astra_db}"]'
@@ -352,6 +360,10 @@ async def __aexit__(
             traceback=traceback,
         )
 
+    @property
+    def namespace(self) -> str:
+        return self._astra_db.namespace
+
     def copy(self) -> AsyncDatabase:
         return AsyncDatabase(
             **self._constructor_params,

diff --git a/tests/idiomatic/unit/test_databases_async.py b/tests/idiomatic/unit/test_databases_async.py
@@ -100,6 +100,11 @@ async def test_database_get_collection_async(
         collection = await async_database.get_collection(TEST_COLLECTION_INSTANCE_NAME)
         assert collection == async_collection_instance
 
+        assert (
+            await getattr(async_database, TEST_COLLECTION_INSTANCE_NAME) == collection
+        )
+        assert await async_database[TEST_COLLECTION_INSTANCE_NAME] == collection
+
         NAMESPACE_2 = "other_namespace"
         collection_ns2 = await async_database.get_collection(
             TEST_COLLECTION_INSTANCE_NAME, namespace=NAMESPACE_2

diff --git a/tests/idiomatic/unit/test_databases_sync.py b/tests/idiomatic/unit/test_databases_sync.py
@@ -101,6 +101,9 @@ def test_database_get_collection_sync(
         collection = sync_database.get_collection(TEST_COLLECTION_INSTANCE_NAME)
         assert collection == sync_collection_instance
 
+        assert getattr(sync_database, TEST_COLLECTION_INSTANCE_NAME) == collection
+        assert sync_database[TEST_COLLECTION_INSTANCE_NAME] == collection
+
         NAMESPACE_2 = "other_namespace"
         collection_ns2 = sync_database.get_collection(
             TEST_COLLECTION_INSTANCE_NAME, namespace=NAMESPACE_2