add delete collection, refactor

chroma_functions.py

@@ -1,4 +1,5 @@
 import chromadb
+from chromadb.errors import NotFoundError
 from typing import List
 
 
@@ -15,6 +16,7 @@ def list_collections() -> List[str]:
 def create_collection(name: str) -> str:
     client = chromadb.PersistentClient()
     collection = client.create_collection(name=name)
+
     return getattr(collection, "name", name)
 
 
@@ -24,7 +26,11 @@ def delete_collection(name: str) -> None:
 
 
 def count_collection(name: str) -> int:
-    # TODO Implement this method.
-    # The function must use the count method and return how many records are
-    # in the collection. It must handle non-existent collections.
-    raise NotImplemented()
+    client = chromadb.PersistentClient()
+
+    try:
+        collection = client.get_collection(name=name)
+    except NotFoundError:
+        raise
+
+    return collection.count()

cli_parser.py

@@ -25,6 +25,11 @@ def build_parser() -> argparse.ArgumentParser:
     )
     delete_parser.add_argument("name", help="Name of the collection to delete.")
 
-    # TODO add "count <name>" command.
+    count_parser = subparsers.add_parser(
+        "count",
+        aliases=["co"],
+        help="Count records in a collection from the local Chroma database.",
+    )
+    count_parser.add_argument("name", help="Name of the collection to count.")
 
     return parser

main.py

@@ -1,7 +1,12 @@
 from __future__ import annotations
 
 from chromadb.errors import NotFoundError, InternalError
-from chroma_functions import create_collection, delete_collection, list_collections
+from chroma_functions import (
+    count_collection,
+    create_collection,
+    delete_collection,
+    list_collections,
+)
 
 from cli_parser import build_parser
 
@@ -17,16 +22,18 @@ def main() -> int:
 
         for name in collections:
             print(name)
+
         return 0
 
     if args.command in {"create-collection", "cc"}:
         try:
-            create_collection(args.name)
+            collection = create_collection(args.name)
         except InternalError:
             print(f"Collection '{args.name}' already exists.")
             return 1
 
-        print(f"Created collection '{args.name}'.")
+        print(f"Created collection '{collection}'.")
+
         return 0
 
     if args.command in {"delete-collection", "dc"}:
@@ -37,11 +44,22 @@ def main() -> int:
             return 1
 
         print(f"Deleted collection '{args.name}'.")
+
         return 0
 
-    # TODO Implement the count command.
+    if args.command in {"count", "co"}:
+        try:
+            count = count_collection(args.name)
+        except NotFoundError:
+            print(f"Collection '{args.name}' does not exist.")
+            return 1
+
+        print(count)
+
+        return 0
 
     print("Nothing to do. Use -h to see available commands.")
+
     return 0