Stubs for google.cloud.ndb the Google Cloud Datastore ndb client library

pyrightconfig.stricter.json

@@ -50,6 +50,7 @@
         "stubs/paramiko",
         "stubs/prettytable",
         "stubs/protobuf",
+        "stubs/google-cloud-ndb",
         "stubs/psutil",
         "stubs/psycopg2",
         "stubs/Pygments",

stubs/google-cloud-ndb/@tests/stubtest_allowlist.txt

@@ -0,0 +1,4 @@
+# inconsistency of signatures between stub and implementation (cls vs self)
+google.cloud.ndb.ModelAdapter.__new__
+google.cloud.ndb.metadata.EntityGroup.__new__
+google.cloud.ndb.model.ModelAdapter.__new__

stubs/google-cloud-ndb/METADATA.toml

@@ -0,0 +1,2 @@
+version = "1.9"
+requires = ["types-six"]

stubs/google-cloud-ndb/google/cloud/ndb/__init__.pyi

@@ -0,0 +1,109 @@
+from google.cloud.ndb._datastore_api import EVENTUAL as EVENTUAL, EVENTUAL_CONSISTENCY as EVENTUAL_CONSISTENCY, STRONG as STRONG
+from google.cloud.ndb._datastore_query import Cursor as Cursor, QueryIterator as QueryIterator
+from google.cloud.ndb._transaction import (
+    in_transaction as in_transaction,
+    non_transactional as non_transactional,
+    transaction as transaction,
+    transaction_async as transaction_async,
+    transactional as transactional,
+    transactional_async as transactional_async,
+    transactional_tasklet as transactional_tasklet,
+)
+from google.cloud.ndb.client import Client as Client
+from google.cloud.ndb.context import (
+    AutoBatcher as AutoBatcher,
+    Context as Context,
+    ContextOptions as ContextOptions,
+    TransactionOptions as TransactionOptions,
+    get_context as get_context,
+    get_toplevel_context as get_toplevel_context,
+)
+from google.cloud.ndb.global_cache import GlobalCache as GlobalCache, MemcacheCache as MemcacheCache, RedisCache as RedisCache
+from google.cloud.ndb.key import Key as Key
+from google.cloud.ndb.model import (
+    BadProjectionError as BadProjectionError,
+    BlobKey as BlobKey,
+    BlobKeyProperty as BlobKeyProperty,
+    BlobProperty as BlobProperty,
+    BooleanProperty as BooleanProperty,
+    ComputedProperty as ComputedProperty,
+    ComputedPropertyError as ComputedPropertyError,
+    DateProperty as DateProperty,
+    DateTimeProperty as DateTimeProperty,
+    Expando as Expando,
+    FloatProperty as FloatProperty,
+    GenericProperty as GenericProperty,
+    GeoPt as GeoPt,
+    GeoPtProperty as GeoPtProperty,
+    Index as Index,
+    IndexProperty as IndexProperty,
+    IndexState as IndexState,
+    IntegerProperty as IntegerProperty,
+    InvalidPropertyError as InvalidPropertyError,
+    JsonProperty as JsonProperty,
+    KeyProperty as KeyProperty,
+    KindError as KindError,
+    LocalStructuredProperty as LocalStructuredProperty,
+    MetaModel as MetaModel,
+    Model as Model,
+    ModelAdapter as ModelAdapter,
+    ModelAttribute as ModelAttribute,
+    ModelKey as ModelKey,
+    PickleProperty as PickleProperty,
+    Property as Property,
+    ReadonlyPropertyError as ReadonlyPropertyError,
+    Rollback as Rollback,
+    StringProperty as StringProperty,
+    StructuredProperty as StructuredProperty,
+    TextProperty as TextProperty,
+    TimeProperty as TimeProperty,
+    UnprojectedPropertyError as UnprojectedPropertyError,
+    User as User,
+    UserNotFoundError as UserNotFoundError,
+    UserProperty as UserProperty,
+    delete_multi as delete_multi,
+    delete_multi_async as delete_multi_async,
+    get_indexes as get_indexes,
+    get_indexes_async as get_indexes_async,
+    get_multi as get_multi,
+    get_multi_async as get_multi_async,
+    make_connection as make_connection,
+    put_multi as put_multi,
+    put_multi_async as put_multi_async,
+)
+from google.cloud.ndb.polymodel import PolyModel as PolyModel
+from google.cloud.ndb.query import (
+    AND as AND,
+    OR as OR,
+    ConjunctionNode as ConjunctionNode,
+    DisjunctionNode as DisjunctionNode,
+    FalseNode as FalseNode,
+    FilterNode as FilterNode,
+    Node as Node,
+    Parameter as Parameter,
+    ParameterizedFunction as ParameterizedFunction,
+    ParameterizedThing as ParameterizedThing,
+    ParameterNode as ParameterNode,
+    PostFilterNode as PostFilterNode,
+    Query as Query,
+    QueryOptions as QueryOptions,
+    RepeatedStructuredPropertyPredicate as RepeatedStructuredPropertyPredicate,
+    gql as gql,
+)
+from google.cloud.ndb.tasklets import (
+    Future as Future,
+    QueueFuture as QueueFuture,
+    ReducingFuture as ReducingFuture,
+    Return as Return,
+    SerialQueueFuture as SerialQueueFuture,
+    add_flow_exception as add_flow_exception,
+    make_context as make_context,
+    make_default_context as make_default_context,
+    set_context as set_context,
+    sleep as sleep,
+    synctasklet as synctasklet,
+    tasklet as tasklet,
+    toplevel as toplevel,
+    wait_all as wait_all,
+    wait_any as wait_any,
+)

stubs/google-cloud-ndb/google/cloud/ndb/_batch.pyi

@@ -0,0 +1,3 @@
+from typing import Any
+
+def get_batch(batch_cls, options: Any | None = ...): ...

stubs/google-cloud-ndb/google/cloud/ndb/_cache.pyi

@@ -0,0 +1,72 @@
+from typing import Any
+
+from google.cloud.ndb import tasklets as tasklets
+
+class ContextCache:
+    def get_and_validate(self, key): ...
+
+class _GlobalCacheBatch:
+    def full(self): ...
+    def idle_callback(self) -> None: ...
+    def done_callback(self, cache_call) -> None: ...
+    def make_call(self) -> None: ...
+    def future_info(self, key) -> None: ...
+
+global_get: Any
+
+class _GlobalCacheGetBatch(_GlobalCacheBatch):
+    todo: Any
+    keys: Any
+    def __init__(self, ignore_options) -> None: ...
+    def add(self, key): ...
+    def done_callback(self, cache_call) -> None: ...
+    def make_call(self): ...
+    def future_info(self, key): ...
+
+def global_set(key, value, expires: Any | None = ..., read: bool = ...): ...
+
+class _GlobalCacheSetBatch(_GlobalCacheBatch):
+    expires: Any
+    todo: object
+    futures: object
+    def __init__(self, options) -> None: ...
+    def done_callback(self, cache_call) -> None: ...
+    def add(self, key, value): ...
+    def make_call(self): ...
+    def future_info(self, key, value): ...
+
+class _GlobalCacheSetIfNotExistsBatch(_GlobalCacheSetBatch):
+    def add(self, key, value): ...
+    def make_call(self): ...
+    def future_info(self, key, value): ...
+
+global_delete: Any
+
+class _GlobalCacheDeleteBatch(_GlobalCacheBatch):
+    keys: Any
+    futures: Any
+    def __init__(self, ignore_options) -> None: ...
+    def add(self, key): ...
+    def make_call(self): ...
+    def future_info(self, key): ...
+
+global_watch: Any
+
+class _GlobalCacheWatchBatch(_GlobalCacheDeleteBatch):
+    def make_call(self): ...
+    def future_info(self, key): ...
+
+def global_unwatch(key): ...
+
+class _GlobalCacheUnwatchBatch(_GlobalCacheDeleteBatch):
+    def make_call(self): ...
+    def future_info(self, key): ...
+
+global_compare_and_swap: Any
+
+class _GlobalCacheCompareAndSwapBatch(_GlobalCacheSetBatch):
+    def make_call(self): ...
+    def future_info(self, key, value): ...
+
+def is_locked_value(value): ...
+def global_cache_key(key): ...

stubs/google-cloud-ndb/google/cloud/ndb/_datastore_api.pyi

@@ -0,0 +1,5 @@
+from typing_extensions import Literal
+
+EVENTUAL: Literal[2]
+EVENTUAL_CONSISTENCY: Literal[2]
+STRONG: Literal[1]

stubs/google-cloud-ndb/google/cloud/ndb/_datastore_query.pyi

@@ -0,0 +1,22 @@
+from typing import Any
+
+class QueryIterator:
+    def __iter__(self): ...
+    def has_next(self) -> None: ...
+    def has_next_async(self) -> None: ...
+    def probably_has_next(self) -> None: ...
+    def next(self) -> None: ...
+    def cursor_before(self) -> None: ...
+    def cursor_after(self) -> None: ...
+    def index_list(self) -> None: ...
+
+class Cursor:
+    @classmethod
+    def from_websafe_string(cls, urlsafe): ...
+    cursor: Any
+    def __init__(self, cursor: Any | None = ..., urlsafe: Any | None = ...) -> None: ...
+    def to_websafe_string(self): ...
+    def urlsafe(self): ...
+    def __eq__(self, other): ...
+    def __ne__(self, other): ...
+    def __hash__(self): ...

stubs/google-cloud-ndb/google/cloud/ndb/_eventloop.pyi

@@ -0,0 +1,26 @@
+from typing import Any, NamedTuple
+
+class _Event(NamedTuple):
+    when: Any
+    callback: Any
+    args: Any
+    kwargs: Any
+
+class EventLoop:
+    current: Any
+    idlers: Any
+    inactive: int
+    queue: Any
+    rpcs: Any
+    rpc_results: Any
+    def __init__(self) -> None: ...
+    def clear(self) -> None: ...
+    def insort_event_right(self, event) -> None: ...
+    def call_soon(self, callback, *args, **kwargs) -> None: ...
+    def queue_call(self, delay, callback, *args, **kwargs) -> None: ...
+    def queue_rpc(self, rpc, callback) -> None: ...
+    def add_idle(self, callback, *args, **kwargs) -> None: ...
+    def run_idle(self): ...
+    def run0(self): ...
+    def run1(self): ...
+    def run(self) -> None: ...

stubs/google-cloud-ndb/google/cloud/ndb/_options.pyi

@@ -0,0 +1,15 @@
+from typing import Any
+
+class Options:
+    @classmethod
+    def options(cls, wrapped): ...
+    @classmethod
+    def slots(cls): ...
+    def __init__(self, config: Any | None = ..., **kwargs) -> None: ...
+    def __eq__(self, other): ...
+    def __ne__(self, other): ...
+    def copy(self, **kwargs): ...
+    def items(self) -> None: ...
+
+class ReadOptions(Options):
+    def __init__(self, config: Any | None = ..., **kwargs) -> None: ...

stubs/google-cloud-ndb/google/cloud/ndb/_transaction.pyi

@@ -0,0 +1,18 @@
+from typing import Any
+
+def in_transaction(): ...
+def transaction(
+    callback, retries=..., read_only: bool = ..., join: bool = ..., xg: bool = ..., propagation: Any | None = ...
+): ...
+def transaction_async(
+    callback, retries=..., read_only: bool = ..., join: bool = ..., xg: bool = ..., propagation: Any | None = ...
+): ...
+def transaction_async_(
+    callback, retries=..., read_only: bool = ..., join: bool = ..., xg: bool = ..., propagation: Any | None = ...
+): ...
+def transactional(retries=..., read_only: bool = ..., join: bool = ..., xg: bool = ..., propagation: Any | None = ...): ...
+def transactional_async(retries=..., read_only: bool = ..., join: bool = ..., xg: bool = ..., propagation: Any | None = ...): ...
+def transactional_tasklet(
+    retries=..., read_only: bool = ..., join: bool = ..., xg: bool = ..., propagation: Any | None = ...
+): ...
+def non_transactional(allow_existing: bool = ...): ...

stubs/google-cloud-ndb/google/cloud/ndb/blobstore.pyi

@@ -0,0 +1,65 @@
+from typing import Any
+
+from google.cloud.ndb import model
+
+BlobKey: Any
+BLOB_INFO_KIND: str
+BLOB_MIGRATION_KIND: str
+BLOB_KEY_HEADER: str
+BLOB_RANGE_HEADER: str
+MAX_BLOB_FETCH_SIZE: int
+UPLOAD_INFO_CREATION_HEADER: str
+BlobKeyProperty = model.BlobKeyProperty
+
+class BlobFetchSizeTooLargeError:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class BlobInfo:
+    def __init__(self, *args, **kwargs) -> None: ...
+    @classmethod
+    def get(cls, *args, **kwargs) -> None: ...
+    @classmethod
+    def get_async(cls, *args, **kwargs) -> None: ...
+    @classmethod
+    def get_multi(cls, *args, **kwargs) -> None: ...
+    @classmethod
+    def get_multi_async(cls, *args, **kwargs) -> None: ...
+
+class BlobInfoParseError:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class BlobNotFoundError:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class BlobReader:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def create_upload_url(*args, **kwargs) -> None: ...
+def create_upload_url_async(*args, **kwargs) -> None: ...
+
+class DataIndexOutOfRangeError:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def delete(*args, **kwargs) -> None: ...
+def delete_async(*args, **kwargs) -> None: ...
+def delete_multi(*args, **kwargs) -> None: ...
+def delete_multi_async(*args, **kwargs) -> None: ...
+
+class Error:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def fetch_data(*args, **kwargs) -> None: ...
+def fetch_data_async(*args, **kwargs) -> None: ...
+
+get: Any
+get_async: Any
+get_multi: Any
+get_multi_async: Any
+
+class InternalError:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def parse_blob_info(*args, **kwargs) -> None: ...
+
+class PermissionDeniedError:
+    def __init__(self, *args, **kwargs) -> None: ...

stubs/google-cloud-ndb/google/cloud/ndb/client.pyi

@@ -0,0 +1,21 @@
+from typing import Any
+
+DATASTORE_API_HOST: Any
+
+class Client:
+    SCOPE: Any
+    namespace: Any
+    host: Any
+    client_info: Any
+    secure: Any
+    stub: Any
+    def __init__(self, project: Any | None = ..., namespace: Any | None = ..., credentials: Any | None = ...) -> None: ...
+    def context(
+        self,
+        namespace=...,
+        cache_policy: Any | None = ...,
+        global_cache: Any | None = ...,
+        global_cache_policy: Any | None = ...,
+        global_cache_timeout_policy: Any | None = ...,
+        legacy_data: bool = ...,
+    ) -> None: ...