dara-core 1.17.6__py3-none-any.whl → 1.18.0__py3-none-any.whl

dara/core/internal/cache_store/lru.py

@@ -88,12 +88,13 @@ class LRUCache(CacheStoreImpl[LruCachePolicy]):
             self.cache.pop(key, None)
             return node.value
 
-    async def get(self, key: str, unpin: bool = False) -> Optional[Any]:
+    async def get(self, key: str, unpin: bool = False, raise_for_missing: bool = False) -> Optional[Any]:
         """
         Retrieve a value from the cache.
 
         :param key: The key of the value to retrieve.
         :param unpin: If true, the entry will be unpinned if it is pinned.
+        :param raise_for_missing: If true, an exception will be raised if the entry is not found
         :return: The value associated with the key, or None if the key is not in the cache.
         """
         async with self.lock:
@@ -103,6 +104,9 @@ class LRUCache(CacheStoreImpl[LruCachePolicy]):
                     node.pin = False
                 self._move_to_front(node)
                 return node.value
+
+        if raise_for_missing:
+            raise KeyError(f'No cache entry found for {key}')
         return None
 
     async def set(self, key: str, value: Any, pin: bool = False) -> None:

dara/core/internal/cache_store/ttl.py

@@ -52,12 +52,13 @@ class TTLCache(CacheStoreImpl[TTLCachePolicy]):
             _, key = heapq.heappop(self.expiration_heap)
             self.unpinned_cache.pop(key, None)
 
-    async def get(self, key: str, unpin: bool = False) -> Any:
+    async def get(self, key: str, unpin: bool = False, raise_for_missing: bool = False) -> Any:
         """
         Retrieve a value from the cache.
 
         :param key: The key of the value to retrieve.
         :param unpin: If true, the entry will be unpinned if it is pinned.
+        :param raise_for_missing: If true, an exception will be raised if the entry is not found
         :return: The value associated with the key, or None if the key is not in the cache.
         """
         async with self.lock:
@@ -75,6 +76,8 @@ class TTLCache(CacheStoreImpl[TTLCachePolicy]):
             elif key in self.unpinned_cache:
                 return self.unpinned_cache[key].value
 
+            if raise_for_missing:
+                raise KeyError(f'No cache entry found for {key}')
             return None
 
     async def set(self, key: str, value: Any, pin: bool = False) -> None:

dara/core/internal/cgroup.py

@@ -22,7 +22,7 @@ import sys
 
 from dara.core.logging import dev_logger, eng_logger
 
-CGROUP_V2_INDICATOR_PATH = '/sys/fs/cgroup/cgroup.controllers'   # Used to determine whether we're using cgroupv2
+CGROUP_V2_INDICATOR_PATH = '/sys/fs/cgroup/cgroup.controllers'  # Used to determine whether we're using cgroupv2
 
 CGROUP_V1_MEM_PATH = '/sys/fs/cgroup/memory/memory.limit_in_bytes'
 CGROUP_V2_MEM_PATH = '/sys/fs/cgroup/memory.max'

dara/core/internal/dependency_resolution.py

@@ -30,20 +30,18 @@ from dara.core.logging import dev_logger
 
 
 class ResolvedDerivedVariable(TypedDict):
-    deps: List[int]
     type: Literal['derived']
     uid: str
     values: List[Any]
-    force: bool
+    force_key: Optional[str]
 
 
 class ResolvedDerivedDataVariable(TypedDict):
-    deps: List[int]
     type: Literal['derived-data']
     uid: str
     values: List[Any]
     filters: Optional[Union[FilterQuery, dict]]
-    force: bool
+    force_key: Optional[str]
 
 
 class ResolvedDataVariable(TypedDict):
@@ -64,7 +62,9 @@ def is_resolved_derived_variable(obj: Any) -> TypeGuard[ResolvedDerivedVariable]
     return isinstance(obj, dict) and 'uid' in obj and obj.get('type') == 'derived'
 
 
-def is_resolved_derived_data_variable(obj: Any) -> TypeGuard[ResolvedDerivedDataVariable]:
+def is_resolved_derived_data_variable(
+    obj: Any,
+) -> TypeGuard[ResolvedDerivedDataVariable]:
     return isinstance(obj, dict) and 'uid' in obj and obj.get('type') == 'derived-data'
 
 
@@ -76,9 +76,29 @@ def is_resolved_switch_variable(obj: Any) -> TypeGuard[ResolvedSwitchVariable]:
     return isinstance(obj, dict) and 'uid' in obj and obj.get('type') == 'switch'
 
 
+def clean_force_key(value: Any) -> Any:
+    """
+    Clean an argument to a value to remove force keys
+    """
+    if value is None:
+        return value
+
+    if isinstance(value, dict):
+        # Remove force key from the value
+        value.pop('force_key', None)
+        return {k: clean_force_key(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [clean_force_key(v) for v in value]
+    return value
+
+
 async def resolve_dependency(
     entry: Union[
-        ResolvedDerivedDataVariable, ResolvedDataVariable, ResolvedDerivedVariable, ResolvedSwitchVariable, Any
+        ResolvedDerivedDataVariable,
+        ResolvedDataVariable,
+        ResolvedDerivedVariable,
+        ResolvedSwitchVariable,
+        Any,
     ],
     store: CacheStore,
     task_mgr: TaskManager,
@@ -127,13 +147,21 @@ async def _resolve_derived_data_var(entry: ResolvedDerivedDataVariable, store: C
     input_values: List[Any] = entry.get('values', [])
 
     result = await DerivedDataVariable.resolve_value(
-        data_var, dv_var, store, task_mgr, input_values, entry.get('force', False), entry.get('filters', None)
+        data_entry=data_var,
+        dv_entry=dv_var,
+        store=store,
+        task_mgr=task_mgr,
+        args=input_values,
+        filters=entry.get('filters', None),
+        force_key=entry.get('force_key'),
     )
     return remove_index(result)
 
 
 async def _resolve_derived_var(
-    derived_variable_entry: ResolvedDerivedVariable, store: CacheStore, task_mgr: TaskManager
+    derived_variable_entry: ResolvedDerivedVariable,
+    store: CacheStore,
+    task_mgr: TaskManager,
 ):
     """
     Resolve a derived variable from the registry and get it's new value based on the dynamic variable mapping passed
@@ -149,7 +177,11 @@ async def _resolve_derived_var(
     var = await registry_mgr.get(derived_variable_registry, str(derived_variable_entry.get('uid')))
     input_values: List[Any] = derived_variable_entry.get('values', [])
     result = await DerivedVariable.get_value(
-        var, store, task_mgr, input_values, derived_variable_entry.get('force', False)
+        var_entry=var,
+        store=store,
+        task_mgr=task_mgr,
+        args=input_values,
+        force_key=derived_variable_entry.get('force_key'),
     )
     return result['value']
 
@@ -228,7 +260,11 @@ def _evaluate_condition(condition: dict) -> bool:
             raise ValueError(f'Unknown condition operator: {operator}')
 
 
-async def _resolve_switch_var(switch_variable_entry: ResolvedSwitchVariable, store: CacheStore, task_mgr: TaskManager):
+async def _resolve_switch_var(
+    switch_variable_entry: ResolvedSwitchVariable,
+    store: CacheStore,
+    task_mgr: TaskManager,
+):
     """
     Resolve a switch variable by evaluating its constituent parts and returning the appropriate value.
 

dara/core/internal/devtools.py

@@ -36,7 +36,7 @@ def print_stacktrace():
     trc = 'Traceback (most recent call last):\n'
     stackstr = trc + ''.join(traceback.format_list(stack))
     if exc is not None:
-        stackstr += '  ' + traceback.format_exc().lstrip(trc)   # pylint:disable=bad-str-strip-call
+        stackstr += '  ' + traceback.format_exc().lstrip(trc)
     else:
         stackstr += '   Exception'
 
@@ -52,7 +52,7 @@ def handle_system_exit(error_msg: str):
     try:
         yield
     except SystemExit as e:
-        raise InterruptedError(error_msg).with_traceback(e.__traceback__)
+        raise InterruptedError(error_msg) from e
 
 
 def get_error_for_channel() -> dict:

dara/core/internal/download.py

@@ -18,7 +18,8 @@ limitations under the License.
 from __future__ import annotations
 
 import os
-from typing import Awaitable, Callable, Optional, Tuple
+from collections.abc import Awaitable
+from typing import Callable, Optional, Tuple
 from uuid import uuid4
 
 import anyio
@@ -37,13 +38,13 @@ class DownloadDataEntry(BaseModel):
     file_path: str
     cleanup_file: bool
     identity_name: Optional[str] = None
-    download: Callable[['DownloadDataEntry'], Awaitable[Tuple[anyio.AsyncFile, Callable[..., Awaitable]]]]
+    download: Callable[[DownloadDataEntry], Awaitable[Tuple[anyio.AsyncFile, Callable[..., Awaitable]]]]
     """Handler for getting the file from the entry"""
 
 
 DownloadRegistryEntry = CachedRegistryEntry(
     uid='_dara_download', cache=Cache.Policy.TTL(ttl=60 * 10)
-)   # expire the codes after 10 minutes
+)  # expire the codes after 10 minutes
 
 
 async def download(data_entry: DownloadDataEntry) -> Tuple[anyio.AsyncFile, Callable[..., Awaitable]]:

dara/core/internal/encoder_registry.py

@@ -14,13 +14,13 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-# pylint: disable=unnecessary-lambda
+
+from collections.abc import MutableMapping
 from inspect import Parameter, isclass
 from typing import (
     Any,
     Callable,
     Dict,
-    MutableMapping,
     Optional,
     Type,
     Union,
@@ -99,11 +99,11 @@ def _tuple_key_deserialize(d):
         if isinstance(key, str) and key.startswith('__tuple__'):
             key_list = []
             for each in key[10:-1].split(', '):
-                if (each.startswith("'") and each.endswith(("'"))) or (each.startswith('"') and each.endswith(('"'))):
+                if (each.startswith("'") and each.endswith("'")) or (each.startswith('"') and each.endswith('"')):
                     key_list.append(each[1:-1])
                 else:
                     key_list.append(each)
-            encoded_key = encoded_key = tuple(key_list)
+            encoded_key = tuple(key_list)
         else:
             encoded_key = key
 
@@ -112,7 +112,7 @@ def _tuple_key_deserialize(d):
     return encoded_dict
 
 
-def _df_deserialize(x):  # pylint: disable=inconsistent-return-statements
+def _df_deserialize(x):
     """
     A function to deserialize data into a DataFrame
 
@@ -240,14 +240,14 @@ def deserialize(value: Any, typ: Optional[Type]):
         return value
 
     # Already matches type
-    if type(value) == typ:
+    if type(value) is typ:
         return value
 
     # Handle Optional[foo] / Union[foo, None] -> call deserialize(value, foo)
     if get_origin(typ) == Union:
         args = get_args(typ)
         if len(args) == 2 and type(None) in args:
-            not_none_arg = args[0] if args[0] != type(None) else args[1]
+            not_none_arg = args[0] if args[0] is not type(None) else args[1]
             return deserialize(value, not_none_arg)
 
     try:

dara/core/internal/execute_action.py

@@ -14,11 +14,13 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Mapping
 from contextvars import ContextVar
-from typing import Any, Callable, Mapping, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import anyio
 
@@ -30,11 +32,7 @@ from dara.core.interactivity.actions import (
     ActionImpl,
 )
 from dara.core.internal.cache_store import CacheStore
-from dara.core.internal.dependency_resolution import (
-    is_resolved_derived_data_variable,
-    is_resolved_derived_variable,
-    resolve_dependency,
-)
+from dara.core.internal.dependency_resolution import resolve_dependency
 from dara.core.internal.encoder_registry import deserialize
 from dara.core.internal.tasks import MetaTask, TaskManager
 from dara.core.internal.utils import run_user_handler
@@ -146,10 +144,6 @@ async def execute_action(
         annotations = action.__annotations__
 
         for key, value in values.items():
-            # Override `force` property to be false
-            if is_resolved_derived_variable(value) or is_resolved_derived_data_variable(value):
-                value['force'] = False
-
             typ = annotations.get(key)
             val = await resolve_dependency(value, store, task_mgr)
             resolved_kwargs[key] = deserialize(val, typ)

dara/core/internal/hashing.py

@@ -31,8 +31,6 @@ def hash_object(obj: Union[BaseModel, dict, None]):
     if isinstance(obj, BaseModel):
         obj = obj.model_dump()
 
-    filter_hash = hashlib.sha1(
-        usedforsecurity=False
-    )   # nosec B303 # we don't use this for security purposes just as a cache key
+    filter_hash = hashlib.sha1(usedforsecurity=False)  # nosec B303 # we don't use this for security purposes just as a cache key
     filter_hash.update(json.dumps(obj or {}, sort_keys=True).encode())
     return filter_hash.hexdigest()

dara/core/internal/import_discovery.py

@@ -93,10 +93,9 @@ def run_discovery(
     # If module root is passed through, use it
     if 'module_root' in kwargs:
         root = kwargs.get('module_root')
-    else:
-        # Try to infer from module_name
-        if module_name is not None:
-            root = module_name.split('.')[0]
+    # Try to infer from module_name
+    elif module_name is not None:
+        root = module_name.split('.')[0]
 
     for k, v in global_symbols.items():
         # Ignore already encountered functions

dara/core/internal/normalization.py

@@ -15,11 +15,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from collections.abc import Mapping
 from typing import (
     Any,
     Generic,
     List,
-    Mapping,
     Optional,
     Tuple,
     TypeVar,
@@ -48,7 +48,7 @@ class Placeholder(TypedDict):
     Placeholder object 'Referrable' objects are replaced with
     """
 
-    __ref: str   # pylint: disable=unused-private-member
+    __ref: str
 
 
 class Referrable(TypedDict):
@@ -56,7 +56,7 @@ class Referrable(TypedDict):
     Describes an object which can be replaced by a Placeholder.
     """
 
-    __typename: str   # pylint: disable=unused-private-member
+    __typename: str
     uid: str
 
 
@@ -133,13 +133,11 @@ def _loop(iterable: JsonLike):
 
 
 @overload
-def normalize(obj: Mapping, check_root: bool = True) -> Tuple[Mapping, Mapping]:
-    ...
+def normalize(obj: Mapping, check_root: bool = True) -> Tuple[Mapping, Mapping]: ...
 
 
 @overload
-def normalize(obj: List, check_root: bool = True) -> Tuple[List, Mapping]:
-    ...
+def normalize(obj: List, check_root: bool = True) -> Tuple[List, Mapping]: ...
 
 
 def normalize(obj: JsonLike, check_root: bool = True) -> Tuple[JsonLike, Mapping]:
@@ -169,7 +167,7 @@ def normalize(obj: JsonLike, check_root: bool = True) -> Tuple[JsonLike, Mapping
         for key, value in _loop(obj):
             # For iterables, recursively call normalize
             if isinstance(value, (dict, list)):
-                _normalized, _lookup = normalize(value)   # type: ignore
+                _normalized, _lookup = normalize(value)  # type: ignore
                 output[key] = _normalized  # type: ignore
                 lookup.update(_lookup)
             else:
@@ -180,13 +178,11 @@ def normalize(obj: JsonLike, check_root: bool = True) -> Tuple[JsonLike, Mapping
 
 
 @overload
-def denormalize(normalized_obj: Mapping, lookup: Mapping) -> Mapping:
-    ...
+def denormalize(normalized_obj: Mapping, lookup: Mapping) -> Mapping: ...
 
 
 @overload
-def denormalize(normalized_obj: List, lookup: Mapping) -> List:
-    ...
+def denormalize(normalized_obj: List, lookup: Mapping) -> List: ...
 
 
 def denormalize(normalized_obj: JsonLike, lookup: Mapping) -> Optional[JsonLike]:
@@ -206,7 +202,7 @@ def denormalize(normalized_obj: JsonLike, lookup: Mapping) -> Optional[JsonLike]
     # Whole object is a placeholder
     if _is_placeholder(normalized_obj):
         ref = normalized_obj['__ref']
-        referrable = lookup[ref] if ref in lookup else None
+        referrable = lookup.get(ref, None)
 
         if isinstance(referrable, (list, dict)):
             return denormalize(referrable, lookup)

dara/core/internal/pandas_utils.py

@@ -15,7 +15,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Optional, TypeVar
+from typing import Optional, TypeVar, cast
 
 from pandas import DataFrame, MultiIndex
 
@@ -31,7 +31,7 @@ def append_index(df: Optional[DataFrame]) -> Optional[DataFrame]:
 
     if INDEX not in df.columns:
         new_df = df.copy()
-        new_df.insert(0, INDEX, range(0, len(df.index)))
+        new_df.insert(0, INDEX, range(0, len(df.index)))  # type: ignore
         return new_df
 
     return df
@@ -47,7 +47,7 @@ def remove_index(value: value_type) -> value_type:
     Otherwise return same value untouched.
     """
     if isinstance(value, DataFrame):
-        return value.drop(columns=['__index__'], inplace=False, errors='ignore')
+        return cast(value_type, value.drop(columns=['__index__'], inplace=False, errors='ignore'))
 
     return value
 

dara/core/internal/pool/task_pool.py

@@ -16,10 +16,11 @@ limitations under the License.
 """
 
 import atexit
+from collections.abc import Coroutine
 from contextlib import contextmanager
 from datetime import datetime
 from multiprocessing import active_children
-from typing import Any, Callable, Coroutine, Dict, Optional, cast
+from typing import Any, Callable, Dict, Optional, Union, cast
 
 import anyio
 from anyio.abc import TaskGroup
@@ -102,16 +103,18 @@ class TaskPool:
             try:
                 await wait_while(
                     lambda: self.status != PoolStatus.RUNNING
-                    or not len(self.workers) == self.desired_workers
+                    or len(self.workers) != self.desired_workers
                     or not all(w.status == WorkerStatus.IDLE for w in self.workers.values()),
                     timeout=timeout,
                 )
-            except TimeoutError:
-                raise RuntimeError('Failed to start pool')
+            except TimeoutError as e:
+                raise RuntimeError('Failed to start pool') from e
         else:
             raise RuntimeError('Pool already started')
 
-    def submit(self, task_uid: str, function_name: str, args: tuple = (), kwargs: dict = {}) -> TaskDefinition:
+    def submit(
+        self, task_uid: str, function_name: str, args: Union[tuple, None] = None, kwargs: Union[dict, None] = None
+    ) -> TaskDefinition:
         """
         Submit a new task to the pool
 
@@ -120,6 +123,10 @@ class TaskPool:
         :param args: list of arguments to pass to the function
         :param kwargs: dict of kwargs to pass to the function
         """
+        if args is None:
+            args = ()
+        if kwargs is None:
+            kwargs = {}
         self._check_pool_state()
 
         # Create a task definition to keep track of its progress
@@ -463,9 +470,8 @@ class TaskPool:
                 )
         elif is_log(worker_msg):
             dev_logger.info(f'Task: {worker_msg["task_uid"]}', {'logs': worker_msg['log']})
-        elif is_progress(worker_msg):
-            if worker_msg['task_uid'] in self._progress_subscribers:
-                await self._progress_subscribers[worker_msg['task_uid']](worker_msg['progress'], worker_msg['message'])
+        elif is_progress(worker_msg) and worker_msg['task_uid'] in self._progress_subscribers:
+            await self._progress_subscribers[worker_msg['task_uid']](worker_msg['progress'], worker_msg['message'])
 
     async def _wait_queue_depletion(self, timeout: Optional[float] = None):
         """
@@ -478,8 +484,8 @@ class TaskPool:
                 condition=lambda: self.status in (PoolStatus.CLOSED, PoolStatus.RUNNING) and len(self.tasks) > 0,
                 timeout=timeout,
             )
-        except TimeoutError:
-            raise TimeoutError('Tasks are still being executed')
+        except TimeoutError as e:
+            raise TimeoutError('Tasks are still being executed') from e
 
     async def _core_loop(self):
         """

dara/core/internal/pool/utils.py

@@ -89,9 +89,7 @@ def read_from_shared_memory(pointer: SharedMemoryPointer) -> Any:
         data = shared_mem.buf[:data_size]
 
         # Unpickle and deepcopy
-        decoded_payload_shared = pickle.loads(
-            shared_mem.buf
-        )   # nosec B301 # we trust the shared memory pointer passed by the pool
+        decoded_payload_shared = pickle.loads(shared_mem.buf)  # nosec B301 # we trust the shared memory pointer passed by the pool
         decoded_payload = copy.deepcopy(decoded_payload_shared)
 
         # Cleanup
@@ -141,8 +139,8 @@ async def stop_process_async(process: BaseProcess, timeout: float = 3):
         try:
             os.kill(process.pid, signal.SIGKILL)
             await wait_while(process.is_alive, timeout)
-        except OSError:
-            raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}')
+        except OSError as e:
+            raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}') from e
 
     # If it's still alive raise an exception
     if process.is_alive():
@@ -171,8 +169,8 @@ def stop_process(process: BaseProcess, timeout: float = 3):
         try:
             os.kill(process.pid, signal.SIGKILL)
             process.join(timeout)
-        except OSError:
-            raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}')
+        except OSError as e:
+            raise RuntimeError(f'Unable to terminate subprocess with PID {process.pid}') from e
 
     # If it's still alive raise an exception
     if process.is_alive():

dara/core/internal/pool/worker.py

@@ -57,7 +57,8 @@ class StdoutLogger:
             self.channel.worker_api.log(self.task_uid, msg)
 
     def flush(self):
-        sys.__stdout__.flush()
+        if sys.__stdout__:
+            sys.__stdout__.flush()
 
 
 def execute_function(func: Callable, args: tuple, kwargs: dict):
@@ -164,7 +165,7 @@ def worker_loop(worker_params: WorkerParameters, channel: Channel):
 
         # Redirect logs via the channel
         stdout_logger = StdoutLogger(task_uid, channel)
-        sys.stdout = stdout_logger   # type: ignore
+        sys.stdout = stdout_logger  # type: ignore
 
         try:
             payload_pointer = task['payload']

dara/core/internal/port_utils.py

@@ -27,7 +27,7 @@ def is_available(host: str, port: int) -> bool:
     """
     try:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.settimeout(2.0)   # timeout in case port is blocked
+            sock.settimeout(2.0)  # timeout in case port is blocked
             return sock.connect_ex((host, port)) != 0
     except BaseException:
         return False

dara/core/internal/registries.py

@@ -15,8 +15,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from collections.abc import Mapping
 from datetime import datetime
-from typing import Any, Callable, Mapping, Set
+from typing import Any, Callable, Set
 
 from dara.core.auth import BaseAuthConfig
 from dara.core.base_definitions import ActionDef, ActionResolverDef, UploadResolverDef
@@ -31,15 +32,16 @@ from dara.core.interactivity.derived_variable import (
     DerivedVariableRegistryEntry,
     LatestValueRegistryEntry,
 )
+from dara.core.internal.download import DownloadDataEntry
 from dara.core.internal.registry import Registry, RegistryType
 from dara.core.internal.websocket import CustomClientMessagePayload
 from dara.core.persistence import BackendStoreEntry
 
-action_def_registry = Registry[ActionDef](RegistryType.ACTION_DEF, CORE_ACTIONS)   # all registered actions
-action_registry = Registry[ActionResolverDef](RegistryType.ACTION)   # functions for actions requiring backend calls
+action_def_registry = Registry[ActionDef](RegistryType.ACTION_DEF, CORE_ACTIONS)  # all registered actions
+action_registry = Registry[ActionResolverDef](RegistryType.ACTION)  # functions for actions requiring backend calls
 upload_resolver_registry = Registry[UploadResolverDef](
     RegistryType.UPLOAD_RESOLVER
-)   # functions for upload resolvers requiring backend calls
+)  # functions for upload resolvers requiring backend calls
 component_registry = Registry[ComponentTypeAnnotation](RegistryType.COMPONENTS, CORE_COMPONENTS)
 config_registry = Registry[EndpointConfiguration](RegistryType.ENDPOINT_CONFIG)
 data_variable_registry = Registry[DataVariableRegistryEntry](RegistryType.DATA_VARIABLE, allow_duplicates=False)
@@ -69,3 +71,6 @@ custom_ws_handlers_registry = Registry[Callable[[str, CustomClientMessagePayload
 
 backend_store_registry = Registry[BackendStoreEntry](RegistryType.BACKEND_STORE, allow_duplicates=False)
 """map of store uid -> store instance"""
+
+download_code_registry = Registry[DownloadDataEntry](RegistryType.DOWNLOAD_CODE, allow_duplicates=False)
+"""map of download codes -> download data entry, used only to allow overriding download code behaviour via RegistryLookup"""

dara/core/internal/registry.py

@@ -16,8 +16,9 @@ limitations under the License.
 """
 
 import copy
+from collections.abc import MutableMapping
 from enum import Enum
-from typing import Generic, MutableMapping, Optional, TypeVar
+from typing import Generic, Optional, TypeVar
 
 from dara.core.metrics import CACHE_METRICS_TRACKER, total_size
 
@@ -43,6 +44,7 @@ class RegistryType(str, Enum):
     PENDING_TOKENS = 'Pending tokens'
     CUSTOM_WS_HANDLERS = 'Custom WS handlers'
     BACKEND_STORE = 'Backend Store'
+    DOWNLOAD_CODE = 'Download Code'
 
 
 class Registry(Generic[T]):