Skip to content

Commit f200987

Browse files
committed
more _FutureGroup typing
1 parent a55e119 commit f200987

File tree

2 files changed

+5
-5
lines changed

2 files changed

+5
-5
lines changed

tensorizer/_futuregroup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44

55
class _FutureGroup:
6-
def __init__(self, futures: Sequence[concurrent.futures.Future]):
6+
def __init__(self, futures): # type: (Sequence[_Future]) -> None
77
self.futures = futures
88

99
def cancel(self) -> bool:

tensorizer/serialization.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4201,7 +4201,7 @@ def _prepare_for_write_opaque(
42014201
w.set_min_file_version_number(OPAQUE_TENSORIZER_VERSION)
42024202

42034203
@staticmethod
4204-
def _do_clone(write_spec, dependency: Optional[concurrent.futures.Future]):
4204+
def _do_clone(write_spec, dependency: Optional[_Future]):
42054205
if dependency is not None:
42064206
dependency.result(_TIMEOUT)
42074207
write_spec.tensor = write_spec.tensor.clone().detach()
@@ -4234,7 +4234,7 @@ def _prepare_for_write_encryption(
42344234
]
42354235
)
42364236

4237-
clone_tasks = []
4237+
clone_tasks: List[_Future] = []
42384238
for w in shared_write_specs[1:]:
42394239
clone_tasks.append(
42404240
self._computation_pool.submit(
@@ -4401,7 +4401,7 @@ def compute_sha256(
44014401
for w in write_specs:
44024402
old_tensor_data_task = w.tensor_data_task
44034403

4404-
hash_tasks = []
4404+
hash_tasks: List[_Future] = []
44054405
if w.include_crc32:
44064406
crc32_task = self._computation_pool.submit(
44074407
compute_crc32, w, old_tensor_data_task
@@ -4418,7 +4418,7 @@ def compute_sha256(
44184418
self._jobs.extend(hash_tasks)
44194419

44204420
def _do_encryption(self, write_specs: Sequence[_WriteSpec]) -> None:
4421-
def encrypt(write_spec, dependency: _Future):
4421+
def encrypt(write_spec, dependency: Optional[_Future]):
44224422
if dependency is not None:
44234423
dependency.result(_TIMEOUT)
44244424
try:

0 commit comments

Comments
 (0)