wrap_module

penzai.toolshed.auto_nmap.wrap_module(module: types.ModuleType) → AutoNmapModuleWrapper

Wraps a module to automatically apply named_axes.nmap to callables.

The returned object will have similar attributes to module, including all submodules and all callables. Accessing these on the returned object will return an nmap-ped version of the original module callable, or a wrapped submodule.

Parameters:

module – The module to wrap.

Returns:

A wrapper object similar to module, such that callable attributes are wrapped with named_axes.nmap, and submodules are recursively wrapped with wrap_module.