Source code for yt.fields.local_fields

from collections.abc import Callable
from functools import partial
from typing import Any, TypeVar

from yt.funcs import is_sequence
from yt.utilities.logger import ytLogger as mylog

from .field_info_container import FieldInfoContainer
from .field_plugin_registry import register_field_plugin

# workaround mypy not being comfortable around decorator preserving signatures
# adapted from
# https://github.com/python/mypy/issues/1551#issuecomment-253978622
TFun = TypeVar("TFun", bound=Callable[..., Any])


[docs] class LocalFieldInfoContainer(FieldInfoContainer):
[docs] def add_field( self, name, function, sampling_type, *, force_override=False, **kwargs ): from yt.fields.field_functions import validate_field_function validate_field_function(function) if isinstance(name, str) or not is_sequence(name): # the base method only accepts proper tuple field keys # and is only used internally, while this method is exposed to users # and is documented as usable with single strings as name if sampling_type == "particle": ftype = "all" else: ftype = "gas" name = (ftype, name) # Handle the case where the field has already been added. if not force_override and name in self: mylog.warning( "Field %s already exists. To override use `force_override=True`.", name, ) return super().add_field( name, function, sampling_type, force_override=force_override, **kwargs )
# Empty FieldInfoContainer local_fields = LocalFieldInfoContainer(None, [], None) # we define two handles, essentially pointing to the same function but documented differently # yt.add_field() is meant to be used directly, while yt.derived_field is documented # as a decorator. add_field = local_fields.add_field
[docs] class derived_field: # implement a decorator accepting keyword arguments to be passed down to add_field def __init__(self, **kwargs) -> None: self._kwargs = kwargs def __call__(self, f: Callable) -> Callable: partial(local_fields.add_field, function=f)(**self._kwargs) return f
[docs] @register_field_plugin def setup_local_fields(registry, ftype="gas", slice_info=None): # This is easy. We just update with the contents of the local_fields field # info container, and since they are not mutable in any real way, we are # fine. # Note that we actually don't care about the ftype here. for f in local_fields: registry._show_field_errors.append(f) registry.update(local_fields)