Allow Python Future and Promise classes to behave like generic classes

PiperOrigin-RevId: 823076226
Change-Id: I9edd345d77f2cb9860f0da22168e8e0a21bf44e7

Author: jbms

python/tensorstore/future.cc

@@ -703,13 +703,20 @@ See also:
 Group:
   Asynchronous support
 )";
+
+  static const PyMethodDef future_methods[] = {
+      {"__class_getitem__", Py_GenericAlias, METH_O | METH_CLASS, ""},
+      {nullptr},
+  };
   PyType_Slot slots[] = {
       {Py_tp_doc, const_cast<char*>(doc)},
       {Py_tp_alloc, reinterpret_cast<void*>(&FutureAlloc)},
       {Py_tp_dealloc, reinterpret_cast<void*>(&FutureDealloc)},
       {Py_tp_traverse, reinterpret_cast<void*>(&FutureTraverse)},
       {Py_tp_clear, reinterpret_cast<void*>(&FutureClear)},
       {Py_tp_finalize, reinterpret_cast<void*>(&FutureFinalize)},
+      {Py_tp_methods,
+       const_cast<void*>(reinterpret_cast<const void*>(future_methods))},
       {0, nullptr},
   };
   PyType_Spec spec = {};
@@ -1025,12 +1032,18 @@ See also:
 Group:
   Asynchronous support
 )";
+  static const PyMethodDef promise_methods[] = {
+      {"__class_getitem__", Py_GenericAlias, METH_O | METH_CLASS, ""},
+      {nullptr},
+  };
   PyType_Slot slots[] = {
       {Py_tp_doc, const_cast<char*>(doc)},
       {Py_tp_alloc, reinterpret_cast<void*>(&PromiseAlloc)},
       {Py_tp_dealloc, reinterpret_cast<void*>(&PromiseDealloc)},
       {Py_tp_traverse, reinterpret_cast<void*>(&PromiseTraverse)},
       {Py_tp_clear, reinterpret_cast<void*>(&PromiseClear)},
+      {Py_tp_methods,
+       const_cast<void*>(reinterpret_cast<const void*>(promise_methods))},
       {0, nullptr},
   };
   PyType_Spec spec = {};

python/tensorstore/tests/future_test.py

@@ -18,6 +18,7 @@
 import signal
 import threading
 import time
+import types
 
 import pytest
 import tensorstore as ts
@@ -31,6 +32,11 @@ def test_promise_new():
   assert future.result() == 5
 
 
+def test_type_params():
+  assert ts.Promise[int] == types.GenericAlias(ts.Promise, (int,))
+  assert ts.Future[int] == types.GenericAlias(ts.Future, (int,))
+
+
 def test_promise_result_release_gil():
   promise, future = ts.Promise.new()
   t = threading.Thread(target=future.result)