chalkpy 2.90.1__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/features/_class_property.py

@@ -5,6 +5,8 @@ import functools
 from typing import Any, Callable, List, Type, TypeVar, cast
 
 from chalk._lsp.error_builder import FeatureClassErrorBuilder
+from chalk.features.feature_wrapper import UnresolvedFeature
+from chalk.utils.notebook import is_notebook
 
 T = TypeVar("T")
 V = TypeVar("V")
@@ -54,6 +56,11 @@ def classproperty_support(cls: Type[T]) -> Type[T]:
             if (res := self.__chalk_notebook_feature_expressions__.get(item)) is not None:
                 return res
 
+            # If in notebook, fallback to constructing FQN string instead of raising error
+            if is_notebook():
+                fqn = f"{self.namespace}.{item}"
+                return UnresolvedFeature(fqn)
+
             builder: FeatureClassErrorBuilder = self.__chalk_error_builder__
             builder.invalid_attribute(
                 root_feature_str=self.namespace,

chalk/features/_embedding/embedding.py

@@ -25,6 +25,7 @@ from chalk.utils.collections import ensure_tuple
 SUPPORTED_LOCAL_MODELS = {
     "all-MiniLM-L6-v2",  # https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
     "sample-bert",  # For internal Chalk use
+    "sample-linear-nn",  # For internal Chalk use
 }
 
 # This will eventually be included in SUPPORTED_LOCAL_MODELS

chalk/features/_encoding/converter.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import json
 import types
 import typing
@@ -50,6 +51,7 @@ from chalk.features._encoding.pyarrow import (
     rich_to_pyarrow,
 )
 from chalk.features._encoding.rich import structure_primitive_to_rich, unstructure_rich_to_primitive
+from chalk.features.feature_wrapper import UnresolvedFeature
 from chalk.utils.collections import unwrap_annotated_if_needed, unwrap_optional_and_annotated_if_needed
 from chalk.utils.df_utils import pa_array_to_pl_series
 from chalk.utils.json import JSON, TJSON, is_pyarrow_json_type, pyarrow_json_type
@@ -870,6 +872,79 @@ class PrimitiveFeatureConverter(Generic[_TPrim]):
         else:
             raise TypeError(f"Could not convert the pyarrow dtype {dtype} to a protobuf message")
 
+    @classmethod
+    def convert_pa_field_to_proto_field(cls, field: pa.Field) -> pb.Field:
+        """Convert a PyArrow Field to proto Field."""
+        field_proto = pb.Field(
+            name=field.name, arrow_type=cls.convert_pa_dtype_to_proto_dtype(field.type), nullable=field.nullable
+        )
+
+        if field.metadata:
+            # field.metadata is of types dict[bytes, bytes]
+            for k, v in field.metadata.items():
+                field_proto.metadata[k.decode("utf-8")] = v.decode("utf-8")
+
+        return field_proto
+
+    @classmethod
+    def convert_proto_field_to_pa_field(cls, proto_field: pb.Field) -> pa.Field:
+        """Convert a proto Field to PyArrow Field."""
+        arrow_type = cls.convert_proto_dtype_to_pa_dtype(proto_field.arrow_type)
+
+        # don't have to convert back to dict[bytes, bytes] as can initialize with dict[str, str]
+        metadata = dict(proto_field.metadata) if proto_field.metadata else None
+
+        return pa.field(
+            name=proto_field.name,
+            type=arrow_type,
+            nullable=proto_field.nullable,
+            metadata=metadata,
+        )
+
+    @classmethod
+    def convert_pa_schema_to_proto_schema(cls, schema: pa.Schema) -> pb.Schema:
+        schema_proto = pb.Schema(
+            columns=[cls.convert_pa_field_to_proto_field(field) for field in schema],
+        )
+
+        if schema.metadata:
+            # schema.metadata is of types dict[bytes, bytes]
+            for k, v in schema.metadata.items():
+                schema_proto.metadata[k.decode("utf-8")] = v.decode("utf-8")
+
+        return schema_proto
+
+    @classmethod
+    def convert_proto_schema_to_pa_schema(cls, proto_schema: pb.Schema) -> pa.Schema:
+        fields = [cls.convert_proto_field_to_pa_field(proto_field) for proto_field in proto_schema.columns]
+
+        # don't have to convert back to dict[bytes, bytes] as can initialize with dict[str, str]
+        metadata = dict(proto_schema.metadata) if proto_schema.metadata else None
+
+        return pa.schema(fields, metadata=metadata)
+
+    @staticmethod
+    def convert_arrow_table_to_proto(table: pa.Table | pa.RecordBatch) -> pb.TableParquetBytes:
+        if isinstance(table, pa.RecordBatch):
+            table = pa.Table.from_batches([table])
+        elif isinstance(table, pa.Table):
+            pass
+        else:
+            raise TypeError(f"expected pa.Table or pa.RecordBatch, got {type(table)!r}")
+
+        sink = io.BytesIO()
+        import pyarrow.parquet
+
+        pyarrow.parquet.write_table(table, sink)
+        return pb.TableParquetBytes(encoded_parquet_bytes=sink.getvalue())
+
+    @staticmethod
+    def convert_arrow_table_from_proto(proto: pb.TableParquetBytes) -> pa.Table:
+        import pyarrow.parquet
+
+        pf = pyarrow.parquet.ParquetFile(io.BytesIO(proto.encoded_parquet_bytes))
+        return pyarrow.parquet.read_table(pf)
+
     @staticmethod
     def _serialize_pa_decimal_to_pb(value: Union[pa.Decimal128Scalar, pa.Decimal256Scalar]) -> pb.ScalarValue:
         dec_val = value.as_py()
@@ -1183,8 +1258,14 @@ class FeatureConverter(PrimitiveFeatureConverter[_TPrim], Generic[_TPrim, _TRich
         # because it is also used for error handling inside of `from_rich_to_primitive`.
         self._name = name
         if rich_default != ...:
-            # The missing value strategy doesn't really matter because rich_default is not missing
-            primitive_default = self.from_rich_to_primitive(rich_default, missing_value_strategy="allow")
+            # In notebook environments, UnresolvedFeature may be used as a placeholder
+            # for features that can't be resolved due to a stale registry.
+            # Treat these as missing defaults since they're not concrete values.
+            if isinstance(rich_default, UnresolvedFeature):
+                rich_default = ...
+            else:
+                # The missing value strategy doesn't really matter because rich_default is not missing
+                primitive_default = self.from_rich_to_primitive(rich_default, missing_value_strategy="allow")
         super().__init__(
             name, is_nullable=is_nullable, pyarrow_dtype=pyarrow_dtype, primitive_default=primitive_default
         )

chalk/features/feature_field.py

@@ -4,11 +4,12 @@ import copy
 import dataclasses
 import functools
 import inspect
+import itertools
 import os
 import re
 import weakref
 from collections.abc import Mapping, MutableMapping
-from datetime import timedelta, datetime
+from datetime import datetime, timedelta
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -28,7 +29,6 @@ from typing import (
     cast,
 )
 
-import itertools
 import numpy as np
 import pyarrow as pa
 
@@ -39,11 +39,11 @@ from chalk.features._encoding.converter import FeatureConverter, JSONCodec, TDec
 from chalk.features._encoding.primitive import TPrimitive
 from chalk.features.feature_set import CURRENT_FEATURE_REGISTRY, FeatureRegistryProtocol
 from chalk.features.feature_wrapper import FeatureWrapper, NearestNeighborException
-from chalk.features.filter import Filter, TimeDelta, ClauseJoinWithAndException
+from chalk.features.filter import ClauseJoinWithAndException, Filter, TimeDelta
 from chalk.features.tag import Tags
 from chalk.features.underscore import Underscore
 from chalk.serialization.parsed_annotation import ParsedAnnotation
-from chalk.utils.collections import ensure_tuple, get_unique_item, FrozenOrderedSet, OrderedSet
+from chalk.utils.collections import FrozenOrderedSet, OrderedSet, ensure_tuple, get_unique_item
 from chalk.utils.duration import CHALK_MAX_TIMEDELTA, Duration, parse_chalk_duration
 from chalk.utils.import_utils import get_type_checking_imports
 from chalk.utils.json import JSON, pyarrow_json_type
@@ -54,7 +54,7 @@ from chalk.utils.pydanticutil.pydantic_compat import (
     is_pydantic_basemodel_instance,
     parse_pydantic_model,
 )
-from chalk.utils.string import to_snake_case, oxford_comma_list
+from chalk.utils.string import oxford_comma_list, to_snake_case
 
 if TYPE_CHECKING:
     from google.protobuf.message import Message as ProtobufMessage
@@ -68,9 +68,9 @@ _TPrim = TypeVar("_TPrim", bound=TPrimitive)
 _logger = get_logger(__name__)
 
 from chalk.features.feature_cache_strategy import (
-    CacheStrategy,
-    CacheNullsType,
     CacheDefaultsType,
+    CacheNullsType,
+    CacheStrategy,
     get_cache_strategy_from_cache_settings,
 )
 
@@ -252,6 +252,7 @@ class Feature(Generic[_TPrim, _TRich]):
         "tags",
         "underlying",
         "underscore_expression",
+        "offline_underscore_expression",
         "version",
         "window_duration",
         "window_durations",
@@ -276,6 +277,7 @@ class Feature(Generic[_TPrim, _TRich]):
         primary: bool | None = None,
         default: _TRich | ellipsis = ...,
         underscore_expression: Underscore | None = None,
+        offline_underscore_expression: Underscore | None = None,
         max_staleness: Duration | None | ellipsis = ...,
         online_store_max_items: int | None = None,
         cache_strategy: CacheStrategy = CacheStrategy.ALL_WITH_BOTH_UNSET,
@@ -386,6 +388,7 @@ class Feature(Generic[_TPrim, _TRich]):
         self._primary = primary
         self._primary_feature: Optional[Feature] = None
         self.underscore_expression: Underscore | None = underscore_expression
+        self.offline_underscore_expression: Underscore | None = offline_underscore_expression
         self.is_distance_pseudofeature = is_distance_pseudofeature
 
         self._raw_max_staleness = max_staleness
@@ -503,7 +506,7 @@ class Feature(Generic[_TPrim, _TRich]):
 
     @property
     def converter(self) -> FeatureConverter:
-        from chalk.features import DataFrame, Vector, Tensor
+        from chalk.features import DataFrame, Tensor, Vector
 
         self._converter_entered += 1
 
@@ -923,7 +926,7 @@ class Feature(Generic[_TPrim, _TRich]):
 
     def __repr__(self):
         try:
-            root_fqn=self.root_fqn
+            root_fqn = self.root_fqn
         except:
             # self.root_fqn is a property, if it failed then just return the object repr
             return object.__repr__(self)
@@ -1367,7 +1370,7 @@ class Feature(Generic[_TPrim, _TRich]):
                 ),
                 label="invalid nearest neighbor join",
                 range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                      or self.lsp_error_builder.property_range(self.attribute_name),
+                or self.lsp_error_builder.property_range(self.attribute_name),
                 code="32",
                 raise_error=TypeError,
             )
@@ -1387,7 +1390,7 @@ class Feature(Generic[_TPrim, _TRich]):
                     ),
                     label="invalid join",
                     range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                          or self.lsp_error_builder.property_range(self.attribute_name),
+                    or self.lsp_error_builder.property_range(self.attribute_name),
                     code="32",
                     raise_error=TypeError,
                 )
@@ -1400,7 +1403,7 @@ class Feature(Generic[_TPrim, _TRich]):
                     ),
                     label="invalid join",
                     range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                          or self.lsp_error_builder.property_range(self.attribute_name),
+                    or self.lsp_error_builder.property_range(self.attribute_name),
                     code="32",
                     raise_error=TypeError,
                 )
@@ -1424,18 +1427,21 @@ class Feature(Generic[_TPrim, _TRich]):
             )
 
         if not self.is_has_many and not self.is_has_one:
-            assert self.features_cls is not None
-            self.lsp_error_builder.add_diagnostic(
-                message=(
-                    f"The attribute '{self.features_cls.__name__}.{self.attribute_name}' "
-                    f"has a join filter ({join}) but its type annotation '{self.typ}' is not a feature class or DataFrame."
-                ),
-                label="invalid join",
-                range=self.lsp_error_builder.property_value_range(self.attribute_name)
-                or self.lsp_error_builder.property_range(self.attribute_name),
-                code="37",
-                raise_error=TypeError,
-            )
+            # Check if user tried to use DataFrame (even if validation failed)
+            # Use is_dataframe_annotation() to detect DataFrame types without triggering validation errors
+            if not self.typ.is_dataframe_annotation():
+                assert self.features_cls is not None
+                self.lsp_error_builder.add_diagnostic(
+                    message=(
+                        f"The attribute '{self.features_cls.__name__}.{self.attribute_name}' "
+                        f"has a join filter ({join}) but its type annotation '{self.typ}' is not a feature class or DataFrame that links to another feature class."
+                    ),
+                    label="invalid join",
+                    range=self.lsp_error_builder.property_value_range(self.attribute_name)
+                    or self.lsp_error_builder.property_range(self.attribute_name),
+                    code="37",
+                    raise_error=TypeError,
+                )
         if self._join_type == "has_one":
             if self.is_has_many:
                 assert self.features_cls is not None
@@ -1748,6 +1754,7 @@ def feature(
     default: Union[_TRich, ellipsis] = ...,
     underscore: Optional[Underscore] = None,  # Deprecated. Prefer `expression`.
     expression: Optional[Underscore] = None,
+    offline_expression: Optional[Underscore] = None,
     offline_ttl: Optional[Union[ellipsis, Duration]] = ...,
     deprecated: bool = False,
     store_online: bool = True,
@@ -1811,6 +1818,8 @@ def feature(
         ...     total: int = feature(expression=_.subtotal + _.tax, default=0)
 
         See more at https://docs.chalk.ai/docs/expression
+    offline_expression
+        Defines an alternate expression to compute the feature during offline queries.
     dtype
         The backing `pyarrow.DataType` for the feature. This parameter can
         be used to control the storage format of data. For example, if you
@@ -2058,10 +2067,7 @@ def feature(
             if not isinstance(value, Feature):  # pyright: ignore[reportUnnecessaryIsInstance]
                 raise ValueError(f"When `versions` is provided, the values must be features, but `{value}` was given.")
 
-    cache_strategy = get_cache_strategy_from_cache_settings(
-        cache_nulls=cache_nulls,
-        cache_defaults=cache_defaults
-    )
+    cache_strategy = get_cache_strategy_from_cache_settings(cache_nulls=cache_nulls, cache_defaults=cache_defaults)
 
     return cast(
         _TRich,
@@ -2108,6 +2114,7 @@ def feature(
             ),
             default=default,
             underscore_expression=expression if expression is not None else underscore,
+            offline_underscore_expression=offline_expression,
             offline_ttl=offline_ttl,
             is_deprecated=deprecated,
             store_online=store_online,
@@ -2173,7 +2180,7 @@ def has_one(f: Callable[[], Any]) -> Any:
 def has_many(
     f: Callable[[], Any],
     max_staleness: Union[Duration, None, ellipsis] = ...,
-    online_store_max_items: int | None = None
+    online_store_max_items: int | None = None,
 ) -> Any:
     """Specify a feature that represents a one-to-many relationship.
 
@@ -2192,6 +2199,7 @@ def has_many(
         The maximum number of items to cache for the joined feature. The
         items in the joined feature aggregate, storing the latest values
         of the joined feature for each primary key in the joined feature.
+
     Examples
     --------
     >>> from chalk.features import DataFrame, features, has_many
@@ -2207,7 +2215,9 @@ def has_many(
     ...         lambda: User.id == Card.user_id
     ...     )
     """
-    return Feature(join=f, max_staleness=max_staleness,online_store_max_items=online_store_max_items, join_type="has_many")
+    return Feature(
+        join=f, max_staleness=max_staleness, online_store_max_items=online_store_max_items, join_type="has_many"
+    )
 
 
 __all__ = (

chalk/features/feature_set_decorator.py

@@ -141,6 +141,7 @@ def features(
         The `cache_nulls` and `cache_defaults` options can be used together on the same feature with the
         following exceptions: if `cache_nulls=False`, then `cache_defaults` cannot be `"evict_defaults"`, and if
         `cache_nulls="evict_defaults"`, then `cache_defaults` cannot be `False`.
+
     Other Parameters
     ----------------
     cls

chalk/features/feature_wrapper.py

@@ -11,6 +11,7 @@ from chalk.features._chalkop import op, Aggregation
 from chalk.features.filter import Filter
 from chalk.serialization.parsed_annotation import ParsedAnnotation
 from chalk.utils.collections import ensure_tuple
+from chalk.utils.notebook import is_notebook
 
 if TYPE_CHECKING:
     from chalk.features.feature_field import Feature
@@ -22,16 +23,43 @@ class NearestNeighborException(ValueError):
     ...
 
 
+class UnresolvedFeature:
+    """Fallback for features that can't be resolved in notebook environments.
+
+    This allows notebooks to work even when the feature registry is stale or incomplete.
+    The server will validate the feature exists when the query is executed.
+    """
+    __slots__ = ("fqn",)
+
+    def __init__(self, fqn: str):
+        self.fqn = fqn
+        super().__init__()
+
+    def __str__(self):
+        return self.fqn
+
+    def __repr__(self):
+        return f"UnresolvedFeature({self.fqn!r})"
+
+    def __hash__(self):
+        return hash(self.fqn)
+
+    def __eq__(self, other: object):
+        if isinstance(other, UnresolvedFeature):
+            return self.fqn == other.fqn
+        return False
+
+
 class _MarkedUnderlyingFeature:
     __slots__ = ("_fn", "_source", "_debug_info")
 
-    def __init__(self, fn: Callable[[], Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation],
+    def __init__(self, fn: Callable[[], Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation | UnresolvedFeature],
                  debug_info: Any = None) -> None:
         super().__init__()
         self._fn = fn
         self._debug_info = debug_info
 
-    def __call__(self, *args: Any, **kwds: Any) -> Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation:
+    def __call__(self, *args: Any, **kwds: Any) -> Feature | Filter | type[DataFrame] | FeatureWrapper | Aggregation | UnresolvedFeature:
         return self._fn()
 
 
@@ -51,7 +79,7 @@ class FeatureWrapper:
         super().__init__()
         self._chalk_underlying = underlying
 
-    def _chalk_get_underlying(self) -> Feature | Aggregation | Filter | type[DataFrame]:
+    def _chalk_get_underlying(self) -> Feature | Aggregation | Filter | type[DataFrame] | UnresolvedFeature:
         if isinstance(self._chalk_underlying, _MarkedUnderlyingFeature):
             self._chalk_underlying = self._chalk_underlying()
         if isinstance(self._chalk_underlying, FeatureWrapper):
@@ -303,6 +331,12 @@ class FeatureWrapper:
                     if f.attribute_name == item:
                         return FeatureWrapper(underlying.copy_with_path(f))
 
+                if is_notebook():
+                    # Construct FQN by preserving the path from the underlying feature
+                    # If underlying has a path, we need to include it in the FQN
+                    fqn = f"{underlying.root_fqn}.{item}"
+                    return UnresolvedFeature(fqn)
+
                 assert underlying.features_cls is not None
                 underlying.features_cls.__chalk_error_builder__.invalid_attribute(
                     root_feature_str=joined_class.namespace,
@@ -314,6 +348,11 @@ class FeatureWrapper:
                 )
                 assert False, "unreachable"
 
+            # If in notebook, fallback to constructing FQN string instead of raising error
+            if is_notebook():
+                fqn = f"{underlying.fqn}.{item}"
+                return UnresolvedFeature(fqn)
+
             assert underlying.features_cls is not None
             underlying.features_cls.__chalk_error_builder__.invalid_attribute(
                 root_feature_str=underlying.fqn,

chalk/features/hooks.py

@@ -4,12 +4,15 @@ import asyncio
 import inspect
 import logging
 import time  # Added for measuring duration
-from typing import Any, Callable, Iterable, Optional, Set, Tuple, TypeAlias, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, Set, Tuple, TypeAlias, Union, overload
 
 from chalk.features.tag import Environments
 from chalk.utils.collections import ensure_tuple
 from chalk.utils.log_with_context import get_logger
 
+if TYPE_CHECKING:
+    from chalk.features.resolver import ResourceHint
+
 HookFn: TypeAlias = Callable[[], Any]
 
 
@@ -43,21 +46,31 @@ class Hook:
     venv: Optional[str]
     fn: HookFn
     filename: str
+    resource_hint: Optional["ResourceHint"]
+    resource_group: Optional[str]
 
     def __init__(
-        self, fn: HookFn, filename: str, environment: Optional[Environments] = None, venv: Optional[str] = None
+        self,
+        fn: HookFn,
+        filename: str,
+        environment: Optional[Environments] = None,
+        venv: Optional[str] = None,
+        resource_hint: Optional["ResourceHint"] = None,
+        resource_group: Optional[str] = None,
     ):
         super().__init__()
         self.fn = fn
         self.filename = filename
         self.environment = None if environment is None else ensure_tuple(environment)
         self.venv = venv
+        self.resource_hint = resource_hint
+        self.resource_group = resource_group
 
     def __call__(self):
         return self.fn()
 
     def __repr__(self):
-        return f'Hook(filename={self.filename}, fn={self.fn.__name__}", environment={str(self.environment)}, venv={self.venv})'
+        return f'Hook(filename={self.filename}, fn={self.fn.__name__}", environment={str(self.environment)}, venv={self.venv}, resource_hint={self.resource_hint}, resource_group={self.resource_group})'
 
     @classmethod
     async def async_run_all_before_all(cls, environment: str, venv: Optional[str] = None) -> None:
@@ -75,17 +88,46 @@ def before_all(fn: HookFn, /) -> Hook:
 
 @overload
 def before_all(
-    fn: None = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: HookFn,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
+) -> Hook:
+    ...
+
+
+@overload
+def before_all(
+    fn: None = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Callable[[HookFn], Hook]:
     ...
 
 
 def before_all(
-    fn: Optional[HookFn] = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: Optional[HookFn] = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Union[Hook, Callable[[HookFn], Hook]]:
     def decorator(f: HookFn):
         caller_filename = inspect.getsourcefile(f) or "unknown_file"
-        hook = Hook(fn=f, filename=caller_filename, environment=environment, venv=venv)
+        hook = Hook(
+            fn=f,
+            filename=caller_filename,
+            environment=environment,
+            venv=venv,
+            resource_hint=resource_hint,
+            resource_group=resource_group,
+        )
         Hook.before_all.add(hook)
         return hook
 
@@ -93,23 +135,52 @@ def before_all(
 
 
 @overload
-def after_all(fn: HookFn, /, environment: Optional[Environments] = None, venv: Optional[str] = None) -> Hook:
+def after_all(fn: HookFn, /) -> Hook:
+    ...
+
+
+@overload
+def after_all(
+    fn: HookFn,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
+) -> Hook:
     ...
 
 
 @overload
 def after_all(
-    fn: None = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: None = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Callable[[HookFn], Hook]:
     ...
 
 
 def after_all(
-    fn: Optional[HookFn] = None, /, environment: Optional[Environments] = None, venv: Optional[str] = None
+    fn: Optional[HookFn] = None,
+    /,
+    environment: Optional[Environments] = None,
+    venv: Optional[str] = None,
+    resource_hint: Optional["ResourceHint"] = None,
+    resource_group: Optional[str] = None,
 ) -> Union[Hook, Callable[[HookFn], Hook]]:
     def decorator(f: HookFn):
         caller_filename = inspect.getsourcefile(f) or "unknown_file"
-        hook = Hook(fn=f, filename=caller_filename, environment=environment, venv=venv)
+        hook = Hook(
+            fn=f,
+            filename=caller_filename,
+            environment=environment,
+            venv=venv,
+            resource_hint=resource_hint,
+            resource_group=resource_group,
+        )
         Hook.after_all.add(hook)
         return hook