cocoindex 0.1.43__cp311-cp311-macosx_11_0_arm64.whl → 0.1.45__cp311-cp311-macosx_11_0_arm64.whl

cocoindex/runtime.py

@@ -5,12 +5,14 @@ manner.
 
 import threading
 import asyncio
-from typing import Coroutine
+from typing import Any, Coroutine
+
+
 class _ExecutionContext:
     _lock: threading.Lock
     _event_loop: asyncio.AbstractEventLoop | None = None
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._lock = threading.Lock()
 
     @property
@@ -19,11 +21,14 @@ class _ExecutionContext:
         with self._lock:
             if self._event_loop is None:
                 self._event_loop = asyncio.new_event_loop()
-                threading.Thread(target=self._event_loop.run_forever, daemon=True).start()
+                threading.Thread(
+                    target=self._event_loop.run_forever, daemon=True
+                ).start()
             return self._event_loop
 
-    def run(self, coro: Coroutine):
+    def run(self, coro: Coroutine[Any, Any, Any]) -> Any:
         """Run a coroutine in the event loop, blocking until it finishes. Return its result."""
         return asyncio.run_coroutine_threadsafe(coro, self.event_loop).result()
 
+
 execution_context = _ExecutionContext()

cocoindex/setting.py

@@ -1,43 +1,55 @@
 """
 Data types for settings of the cocoindex library.
 """
+
 import os
 
 from typing import Callable, Self, Any, overload
 from dataclasses import dataclass
 
-_app_namespace: str = ''
+_app_namespace: str = ""
+
 
 def get_app_namespace(*, trailing_delimiter: str | None = None) -> str:
     """Get the application namespace. Append the `trailing_delimiter` if not empty."""
-    if _app_namespace == '' or trailing_delimiter is None:
+    if _app_namespace == "" or trailing_delimiter is None:
         return _app_namespace
-    return f'{_app_namespace}{trailing_delimiter}'
+    return f"{_app_namespace}{trailing_delimiter}"
+
 
 def split_app_namespace(full_name: str, delimiter: str) -> tuple[str, str]:
     """Split the full name into the application namespace and the rest."""
     parts = full_name.split(delimiter, 1)
     if len(parts) == 1:
-        return '', parts[0]
+        return "", parts[0]
     return (parts[0], parts[1])
 
-def set_app_namespace(app_namespace: str):
+
+def set_app_namespace(app_namespace: str) -> None:
     """Set the application namespace."""
     global _app_namespace  # pylint: disable=global-statement
     _app_namespace = app_namespace
 
+
 @dataclass
 class DatabaseConnectionSpec:
     """
     Connection spec for relational database.
     Used by both internal and target storage.
     """
+
     url: str
     user: str | None = None
     password: str | None = None
 
-def _load_field(target: dict[str, Any], name: str, env_name: str, required: bool = False,
-                parse: Callable[[str], Any] | None = None):
+
+def _load_field(
+    target: dict[str, Any],
+    name: str,
+    env_name: str,
+    required: bool = False,
+    parse: Callable[[str], Any] | None = None,
+) -> None:
     value = os.getenv(env_name)
     if value is None:
         if required:
@@ -45,9 +57,11 @@ def _load_field(target: dict[str, Any], name: str, env_name: str, required: bool
     else:
         target[name] = value if parse is None else parse(value)
 
+
 @dataclass
 class Settings:
     """Settings for the cocoindex library."""
+
     database: DatabaseConnectionSpec
     app_namespace: str = ""
 
@@ -61,10 +75,11 @@ class Settings:
         _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
         database = DatabaseConnectionSpec(**db_kwargs)
 
-        app_namespace = os.getenv("COCOINDEX_APP_NAMESPACE", '')
+        app_namespace = os.getenv("COCOINDEX_APP_NAMESPACE", "")
 
         return cls(database=database, app_namespace=app_namespace)
 
+
 @dataclass
 class ServerSettings:
     """Settings for the cocoindex server."""
@@ -80,8 +95,12 @@ class ServerSettings:
         """Load settings from environment variables."""
         kwargs: dict[str, Any] = dict()
         _load_field(kwargs, "address", "COCOINDEX_SERVER_ADDRESS")
-        _load_field(kwargs, "cors_origins", "COCOINDEX_SERVER_CORS_ORIGINS",
-                    parse=ServerSettings.parse_cors_origins)
+        _load_field(
+            kwargs,
+            "cors_origins",
+            "COCOINDEX_SERVER_CORS_ORIGINS",
+            parse=ServerSettings.parse_cors_origins,
+        )
         return cls(**kwargs)
 
     @overload
@@ -93,8 +112,12 @@ class ServerSettings:
     def parse_cors_origins(s: str | None) -> list[str] | None: ...
 
     @staticmethod
-    def parse_cors_origins(s):
+    def parse_cors_origins(s: str | None) -> list[str] | None:
         """
         Parse the CORS origins from a string.
         """
-        return [o for e in s.split(",") if (o := e.strip()) != ""] if s is not None else None
+        return (
+            [o for e in s.split(",") if (o := e.strip()) != ""]
+            if s is not None
+            else None
+        )

cocoindex/setup.py

@@ -1,22 +1,26 @@
 from . import flow
 from . import setting
-from . import _engine
+from . import _engine  # type: ignore
+
 
 def sync_setup() -> _engine.SetupStatus:
     flow.ensure_all_flows_built()
     return _engine.sync_setup()
 
+
 def drop_setup(flow_names: list[str]) -> _engine.SetupStatus:
     flow.ensure_all_flows_built()
     return _engine.drop_setup([flow.get_full_flow_name(name) for name in flow_names])
 
+
 def flow_names_with_setup() -> list[str]:
     result = []
     for name in _engine.flow_names_with_setup():
-        app_namespace, name = setting.split_app_namespace(name, '.')
+        app_namespace, name = setting.split_app_namespace(name, ".")
         if app_namespace == setting.get_app_namespace():
             result.append(name)
     return result
 
-def apply_setup_changes(setup_status: _engine.SetupStatus):
+
+def apply_setup_changes(setup_status: _engine.SetupStatus) -> None:
     _engine.apply_setup_changes(setup_status)

cocoindex/sources.py

@@ -1,7 +1,9 @@
 """All builtin sources."""
+
 from . import op
 import datetime
 
+
 class LocalFile(op.SourceSpec):
     """Import data from local file system."""
 
@@ -40,4 +42,4 @@ class AmazonS3(op.SourceSpec):
     binary: bool = False
     included_patterns: list[str] | None = None
     excluded_patterns: list[str] | None = None
-    sqs_queue_url: str | None = None
+    sqs_queue_url: str | None = None

cocoindex/storages.py

@@ -1,4 +1,5 @@
 """All builtin storages."""
+
 from dataclasses import dataclass
 from typing import Sequence
 
@@ -7,11 +8,14 @@ from . import index
 from .auth_registry import AuthEntryReference
 from .setting import DatabaseConnectionSpec
 
+
 class Postgres(op.StorageSpec):
     """Storage powered by Postgres and pgvector."""
+
     database: AuthEntryReference[DatabaseConnectionSpec] | None = None
     table_name: str | None = None
 
+
 @dataclass
 class Qdrant(op.StorageSpec):
     """Storage powered by Qdrant - https://qdrant.tech/."""
@@ -20,61 +24,77 @@ class Qdrant(op.StorageSpec):
     grpc_url: str = "http://localhost:6334/"
     api_key: str | None = None
 
-@dataclass
-class Neo4jConnection:
-    """Connection spec for Neo4j."""
-    uri: str
-    user: str
-    password: str
-    db: str | None = None
 
 @dataclass
 class TargetFieldMapping:
     """Mapping for a graph element (node or relationship) field."""
+
     source: str
     # Field name for the node in the Knowledge Graph.
     # If unspecified, it's the same as `field_name`.
     target: str | None = None
 
+
 @dataclass
 class NodeFromFields:
     """Spec for a referenced graph node, usually as part of a relationship."""
+
     label: str
     fields: list[TargetFieldMapping]
 
+
 @dataclass
 class ReferencedNode:
     """Storage spec for a graph node."""
+
     label: str
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[index.VectorIndexDef] = ()
 
+
 @dataclass
 class Nodes:
     """Spec to map a row to a graph node."""
+
     kind = "Node"
 
     label: str
 
+
 @dataclass
 class Relationships:
     """Spec to map a row to a graph relationship."""
+
     kind = "Relationship"
 
     rel_type: str
     source: NodeFromFields
     target: NodeFromFields
 
+
 # For backwards compatibility only
 NodeMapping = Nodes
 RelationshipMapping = Relationships
 NodeReferenceMapping = NodeFromFields
 
+
+@dataclass
+class Neo4jConnection:
+    """Connection spec for Neo4j."""
+
+    uri: str
+    user: str
+    password: str
+    db: str | None = None
+
+
 class Neo4j(op.StorageSpec):
     """Graph storage powered by Neo4j."""
+
     connection: AuthEntryReference[Neo4jConnection]
     mapping: Nodes | Relationships
 
+
 class Neo4jDeclaration(op.DeclarationSpec):
     """Declarations for Neo4j."""
 
@@ -83,3 +103,26 @@ class Neo4jDeclaration(op.DeclarationSpec):
     nodes_label: str
     primary_key_fields: Sequence[str]
     vector_indexes: Sequence[index.VectorIndexDef] = ()
+
+
+@dataclass
+class KuzuConnection:
+    """Connection spec for Kuzu."""
+
+    api_server_url: str
+
+
+class Kuzu(op.StorageSpec):
+    """Graph storage powered by Kuzu."""
+
+    connection: AuthEntryReference[KuzuConnection]
+    mapping: Nodes | Relationships
+
+
+class KuzuDeclaration(op.DeclarationSpec):
+    """Declarations for Kuzu."""
+
+    kind = "Kuzu"
+    connection: AuthEntryReference[KuzuConnection]
+    nodes_label: str
+    primary_key_fields: Sequence[str]