From 3edc21c66be959f174c522159025867cd1201d3f Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Wed, 24 Apr 2024 09:43:00 -0500
Subject: [PATCH] Add types for WriteOncePersistentDicts
---
loopy/codegen/__init__.py | 5 ++++-
loopy/schedule/__init__.py | 7 ++++++-
loopy/target/execution.py | 12 +++++++++---
setup.py | 2 +-
4 files changed, 20 insertions(+), 6 deletions(-)
diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py
index 68c41336c..fdaa88619 100644
--- a/loopy/codegen/__init__.py
+++ b/loopy/codegen/__init__.py
@@ -312,7 +312,10 @@ class CodeGenerationState:
# }}}
-code_gen_cache = WriteOncePersistentDict(
+code_gen_cache: WriteOncePersistentDict[
+ TranslationUnit,
+ CodeGenerationResult
+] = WriteOncePersistentDict(
"loopy-code-gen-cache-v3-"+DATA_MODEL_VERSION,
key_builder=LoopyKeyBuilder())
diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py
index 5be848e02..6a7339149 100644
--- a/loopy/schedule/__init__.py
+++ b/loopy/schedule/__init__.py
@@ -42,6 +42,8 @@ from loopy.version import DATA_MODEL_VERSION
if TYPE_CHECKING:
from loopy.kernel import LoopKernel
+ from loopy.translation_unit import CallablesTable, TranslationUnit
+
logger = logging.getLogger(__name__)
@@ -2197,7 +2199,10 @@ def _generate_loop_schedules_inner(
# }}}
-schedule_cache = WriteOncePersistentDict(
+schedule_cache: WriteOncePersistentDict[
+ Tuple[LoopKernel, CallablesTable],
+ LoopKernel
+] = WriteOncePersistentDict(
"loopy-schedule-cache-v4-"+DATA_MODEL_VERSION,
key_builder=LoopyKeyBuilder())
diff --git a/loopy/target/execution.py b/loopy/target/execution.py
index 1e49de938..ad49093e2 100644
--- a/loopy/target/execution.py
+++ b/loopy/target/execution.py
@@ -21,7 +21,7 @@ THE SOFTWARE.
"""
-from typing import (Callable, Tuple, Union, Set, FrozenSet, List, Dict,
+from typing import (Callable, Mapping, Tuple, Union, Set, FrozenSet, List, Dict,
Optional, Sequence, Any)
from dataclasses import dataclass
@@ -721,7 +721,10 @@ class ExecutionWrapperGeneratorBase(ABC):
# }}}
-typed_and_scheduled_cache = WriteOncePersistentDict(
+typed_and_scheduled_cache: WriteOncePersistentDict[
+ Tuple[str, TranslationUnit, Optional[Mapping[str, LoopyType]]],
+ TranslationUnit
+] = WriteOncePersistentDict(
"loopy-typed-and-scheduled-cache-v1-"+DATA_MODEL_VERSION,
key_builder=LoopyKeyBuilder())
@@ -729,7 +732,10 @@ typed_and_scheduled_cache = WriteOncePersistentDict(
caches.append(typed_and_scheduled_cache)
-invoker_cache = WriteOncePersistentDict(
+invoker_cache: WriteOncePersistentDict[
+ Tuple[str, TranslationUnit, str],
+ str
+] = WriteOncePersistentDict(
"loopy-invoker-cache-v10-"+DATA_MODEL_VERSION,
key_builder=LoopyKeyBuilder())
diff --git a/setup.py b/setup.py
index d9b8f6643..17f68968f 100644
--- a/setup.py
+++ b/setup.py
@@ -84,7 +84,7 @@ setup(name="loopy",
python_requires="~=3.8",
install_requires=[
- "pytools>=2023.1.1",
+ "pytools>=2024.1.2",
"pymbolic>=2022.1",
"genpy>=2016.1.2",
--
GitLab