Skip to content

Commit db2d0b4

Browse files
committed
REF: Change type hints so Buffer.handle can be int, object, or None
1 parent 2af478f commit db2d0b4

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

cuda_core/cuda/core/experimental/_memory.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
import abc
88
import weakref
9-
from typing import Optional, Tuple, TypeVar
9+
from typing import Optional, Tuple, TypeVar, Union
1010

1111
from cuda.core.experimental._dlpack import DLDeviceType, make_py_capsule
1212
from cuda.core.experimental._stream import default_stream
@@ -18,6 +18,9 @@
1818
# TODO: define a memory property mixin class and make Buffer and
1919
# MemoryResource both inherit from it
2020

21+
DevicePointerT = Union[driver.CUdeviceptr, int, None]
22+
"""A type union of `Cudeviceptr`, `int` and `None` for hinting Buffer.handle."""
23+
2124

2225
class Buffer:
2326
"""Represent a handle to allocated memory.
@@ -81,7 +84,7 @@ def close(self, stream=None):
8184
self._mnff.close(stream)
8285

8386
@property
84-
def handle(self) -> driver.CUdeviceptr:
87+
def handle(self) -> DevicePointerT:
8588
"""Return the buffer handle object.
8689
8790
.. caution::

0 commit comments

Comments
 (0)