handle_local_states

handle_local_states#

penzai.data_effects.local_state.handle_local_states(body: Callable[[Any], Any], category: Hashable | None = None, category_predicate: Callable[[Hashable], bool] | None = None, lazy: Literal[False] = False, state_sharing: Literal['forbidden', 'allowed', 'unsafe'] = 'forbidden', handler_id: str | None = None) tuple[WithFunctionalLocalState, dict[str, Any]][source]#
penzai.data_effects.local_state.handle_local_states(body: Callable[[Any], Any], category: Hashable | None = None, category_predicate: Callable[[Hashable], bool] | None = None, lazy: Literal[True] = False, state_sharing: Literal['forbidden', 'allowed', 'unsafe'] = 'forbidden', handler_id: str | None = None) tuple[WithFunctionalLocalState, Callable[[], dict[str, Any]]]

Extracts local states from a stateful model.

This method the primary way to transform a stateful model into a functional form that can be run.

Parameters:
  • body – A model or submodel with local state.

  • category – The category of states to extract. Not needed if category_predicate is provided.

  • category_predicate – An optional callable that returns True for categories we should take ownership of. Note that states with different categories must still have unique names if they are being handled by the same handler. Not needed if category is provided.

  • lazy – If True, returns a callable that initializes the state, instead of returning the state itself.

  • state_sharing – Strictness for sharing of states. If “forbidden”, state sharing is strictly not allowed. If “allowed”, state sharing is allowed between InitialLocalStateRequest states with identical initializers, and between SharedLocalStateRequest and any other state with the same name. If “unsafe”, any states with the same name will be shared, with the value coming from whichever one was seen last.

  • handler_id – ID to use for the handler. If None, will be inferred.

Returns:

A handler wrapping the model to handle the given states, and an initial state dictionary to pass as the second argument to that handler (or a callable producing that dictionary if lazy was True).