Fix kvstore type annotations and move py.typed file to correct location

PiperOrigin-RevId: 836860557
Change-Id: Ic809c14f2a1ffdba451e05907840581a44e468db

Author: jbms

python/tensorstore/__init__.pyi

@@ -5438,7 +5438,7 @@ class KvStore:
 
         __hash__: typing.ClassVar[None] = None
 
-        def __add__(self, suffix: str) -> KvStore.Spec:
+        def __add__(self, suffix: str | bytes) -> KvStore.Spec:
             """
             Returns a key-value store with the suffix appended to the path.
 
@@ -5507,7 +5507,7 @@ class KvStore:
                 KvStore.Spec({'driver': 'file', 'path': 'tmp/data/'})
             """
 
-        def __truediv__(self, component: str) -> KvStore.Spec:
+        def __truediv__(self, component: str | bytes) -> KvStore.Spec:
             """
             Returns a key-value store with an additional path component joined to the path.
 
@@ -5674,7 +5674,7 @@ class KvStore:
             """
 
         @path.setter
-        def path(self, arg1: str) -> None: ...
+        def path(self, arg1: str | bytes) -> None: ...
 
         @property
         def url(self) -> str:
@@ -5839,7 +5839,7 @@ class KvStore:
           Operators
         """
 
-    def __contains__(self, key: str) -> bool:
+    def __contains__(self, key: str | bytes) -> bool:
         """
         Synchronously checks if the given key is present.
 
@@ -5881,7 +5881,7 @@ class KvStore:
 
     def __copy__(self) -> KvStore: ...
 
-    def __delitem__(self, key: str) -> None:
+    def __delitem__(self, key: str | bytes) -> None:
         """
         Synchronously deletes a single key.
 

python/tensorstore/kvstore.cc

@@ -727,12 +727,12 @@ See also:
 
   cls.def(
       "__contains__",
-      [](const Self& self, std::string_view key) -> bool {
+      [](const Self& self, StrOrBytesView key) -> bool {
         kvstore::ReadOptions options;
         options.byte_range = OptionalByteRangeRequest::Stat();
         auto result = ValueOrThrow(InterruptibleWait([&]() {
           ScopedPyCriticalSection cs(reinterpret_cast<const PyObject*>(&self));
-          return kvstore::Read(self.value, key, std::move(options));
+          return kvstore::Read(self.value, key.value, std::move(options));
         }()));
         return result.has_value();
       },
@@ -904,10 +904,10 @@ See also:
 
   cls.def(
       "__delitem__",
-      [](const Self& self, std::string_view key) {
+      [](const Self& self, StrOrBytesView key) {
         ValueOrThrow(InterruptibleWait([&]() {
           ScopedPyCriticalSection cs(reinterpret_cast<const PyObject*>(&self));
-          return kvstore::Write(self.value, key, std::nullopt);
+          return kvstore::Write(self.value, key.value, std::nullopt);
         }()));
       },
       py::arg("key"), R"(
@@ -1353,9 +1353,9 @@ Modifies a spec.
         ScopedPyCriticalSection cs(reinterpret_cast<const PyObject*>(&self));
         return self.value.path;
       },
-      [](Self& self, std::string_view path) {
+      [](Self& self, StrOrBytesView path) {
         ScopedPyCriticalSection cs(reinterpret_cast<PyObject*>(&self));
-        self.value.path = path;
+        self.value.path = path.value;
       },
       R"(
 Path prefix within the base key-value store.
@@ -1395,10 +1395,10 @@ Path prefix within the base key-value store.
 
   cls.def(
       "__add__",
-      [](const Self& self, std::string_view suffix) -> kvstore::Spec {
+      [](const Self& self, StrOrBytesView  suffix) -> kvstore::Spec {
         ScopedPyCriticalSection cs(reinterpret_cast<const PyObject*>(&self));
         auto new_spec = self.value;
-        new_spec.AppendSuffix(suffix);
+        new_spec.AppendSuffix(suffix.value);
         return new_spec;
       },
       py::arg("suffix"), R"(
@@ -1421,10 +1421,10 @@ ensure there is a :python:`'/'` separator, use :py:obj:`.__truediv__` instead.
 
   cls.def(
       "__truediv__",
-      [](const Self& self, std::string_view component) -> kvstore::Spec {
+      [](const Self& self, StrOrBytesView component) -> kvstore::Spec {
         ScopedPyCriticalSection cs(reinterpret_cast<const PyObject*>(&self));
         auto new_spec = self.value;
-        new_spec.AppendPathComponent(component);
+        new_spec.AppendPathComponent(component.value);
         return new_spec;
       },
       py::arg("component"), R"(