<commit_before># Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
# ==============================================================================
"""Base Haiku module."""

import abc
import contextlib
import functools
import inspect
import re
from typing import (Any, Callable, ContextManager, Dict, Mapping, NamedTuple,
                    Optional, Tuple, Type, TypeVar)

from haiku._src import base
from haiku._src import data_structures
from haiku._src import stateful
from haiku._src import utils
import jax.numpy as jnp

ThreadLocalStack = data_structures.ThreadLocalStack
T = TypeVar("T")
_APPLY_NAME_SCOPE = "__haiku_name_scope"
modules_with_named_call = False


def profiler_name_scopes(enabled=True):
  """Enable/disable profiler name_scopes on all haiku module methods.

  Note: currently only enables for ``__call__``. See: :function:`named_call` if
  you want to annotate other methods explicitly.

  Args:
    enabled: Whether to enable name scopes or not.
  Returns:
    The previous value of the name_scopes setting.
  """
  global modules_with_named_call
  previously_enabled = modules_with_named_call
  modules_with_named_call = enabled
  return previously_enabled


class ModuleMetaclass(abc.ABCMeta):
  """Metaclass for `Module`."""

  def __new__(
      mcs: Type[Type[T]],
      name: str,
      bases: Tuple[Type[Any], ...],
      clsdict: Dict[str, Any],
  ) -> Type[T]:
    method_names = []

    for key, value in clsdict.items():
      if key == "module_name":
        # Don't patch `module_name` in case the user implemented it as a
        # computed property.
        continue

      elif key.startswith("__") and key != "__call__":
        # Don't patch methods like `__getattr__` or `__del__`.
        continue

      elif isinstance(value, property):
        # TODO(tomhennigan) Preserve the type of property subclasses.
        clsdict[key] = property(
            value.fget if not value.fget else wrap_method(key, value.fget),
            value.fset if not value.fset else wrap_method(key, value.fset),
            value.fdel if not value.fdel else wrap_method(key, value.fdel),
            doc=value.__doc__)

      elif inspect.isfunction(value):
        # We defer patching methods until after the type is created such that we
        # can trigger the descriptor binding them to the class.
        method_names.append(key)

    clsdict.setdefault(
        "__repr__",
        lambda module: module._auto_repr)  # pylint: disable=protected-access

    cls = super(ModuleMetaclass, mcs).__new__(mcs, name, bases, clsdict)

    for method_name in method_names:
      # Note: the code below is subtle, we need to ensure that we're wrapping
      # the method bound to the class. In some cases (e.g. `wrapt`) this is
      # important since the method can trigger different behavior when it is
      # bound (e.g. in wrapt `FunctionWrapper.__get__(None, cls)` produces a
      # `BoundFunctionWrapper` which in turn populates the `instance` argument
      # to decorator functions using args[0]).
      # Equivalent to: `cls.__dict__[method_name].__get__(None, cls)`
      method = getattr(cls, method_name)
      method = wrap_method(method_name, method)
      setattr(cls, method_name, method)

    return cls

  def __call__(cls: Type[T], *args, **kwargs) -> T:
    # Call new such that we have an un-initialized module instance that we can
    # still reference even if there is an exception during __init__. This is
    # needed such that we can make sure the name_scope constructed in __init__
    # is closed even if there is an exception.

    # NOTE: We disable pytype since (somewhat surprisingly) this method is bound
    # with the new class and not the metaclass.
    module = cls.__new__(cls, *args, **kwargs)  # pytype: disable=wrong-arg-types

    # Now attempt to initialize the object.
    init = wrap_method("__init__", cls.__init__)
    init(module, *args, **kwargs)

    module._auto_repr = utils.auto_repr(cls, *args, **kwargs)  # pylint: disable=protected-access

    ran_super_ctor = hasattr(module, "module_name")
    if not ran_super_ctor:
      raise ValueError(
          "Constructing an hk.Module without calling the super constructor "
          "is not supported. Add the following as the first line in your "
          "__init__ method:\n\nsuper(%s, self).__init__()" % cls.__name__)

    return module


class MethodContext(NamedTuple):
  r"""Read only state showing the calling context for a method.

  For example lets define two interceptors and print the values in the context.
  Additionally we will make the first interceptor conditionally short circuit,
  since interceptors stack and are run in order, an earlier interceptor can
  decide to call the next interecptor, or short circuit and call the underlying
  method directly:

  >>> module = hk.Linear(1, name="method_context_example")
  >>> short_circuit = False

  >>> def my_interceptor_1(next_fun, args, kwargs, context):
  ...   print('running my_interceptor_1')
  ...   print('- module.name: ', context.module.name)
  ...   print('- method_name: ', context.method_name)
  ...   if short_circuit:
  ...     return context.orig_method(*args, **kwargs)
  ...   else:
  ...     return next_fun(*args, **kwargs)
  >>> def my_interceptor_2(next_fun, args, kwargs, context):
  ...   print('running my_interceptor_2')
  ...   print('- module.name: ', context.module.name)
  ...   print('- method_name: ', context.method_name)
  ...   return next_fun(*args, **kwargs)

  When ``short_circuit=False`` the two interceptors will run in order:

  >>> with hk.experimental.intercept_methods(my_interceptor_1), \
  ...      hk.experimental.intercept_methods(my_interceptor_2):
  ...   _ = module(jnp.ones([1, 1]))
  running my_interceptor_1
  - module.name:  method_context_example
  - method_name:  __call__
  running my_interceptor_2
  - module.name:  method_context_example
  - method_name:  __call__

  Setting ``short_circuit=True`` will cause the first interecptor to call the
  original method (rather than ``next_fun`` which will trigger the next
  interceptor):

  >>> short_circuit = True
  >>> with hk.experimental.intercept_methods(my_interceptor_1), \
  ...      hk.experimental.intercept_methods(my_interceptor_2):
  ...   _ = module(jnp.ones([1, 1]))
  running my_interceptor_1
  - module.name:  method_context_example<commit_msg>Nit: lint MethodContext docsting in Haiku’s module.py<commit_after># Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
# ==============================================================================
"""Base Haiku module."""

import abc
import contextlib
import functools
import inspect
import re
from typing import (Any, Callable, ContextManager, Dict, Mapping, NamedTuple,
                    Optional, Tuple, Type, TypeVar)

from haiku._src import base
from haiku._src import data_structures
from haiku._src import stateful
from haiku._src import utils
import jax.numpy as jnp

ThreadLocalStack = data_structures.ThreadLocalStack
T = TypeVar("T")
_APPLY_NAME_SCOPE = "__haiku_name_scope"
modules_with_named_call = False


def profiler_name_scopes(enabled=True):
  """Enable/disable profiler name_scopes on all haiku module methods.

  Note: currently only enables for ``__call__``. See: :function:`named_call` if
  you want to annotate other methods explicitly.

  Args:
    enabled: Whether to enable name scopes or not.
  Returns:
    The previous value of the name_scopes setting.
  """
  global modules_with_named_call
  previously_enabled = modules_with_named_call
  modules_with_named_call = enabled
  return previously_enabled


class ModuleMetaclass(abc.ABCMeta):
  """Metaclass for `Module`."""

  def __new__(
      mcs: Type[Type[T]],
      name: str,
      bases: Tuple[Type[Any], ...],
      clsdict: Dict[str, Any],
  ) -> Type[T]:
    method_names = []

    for key, value in clsdict.items():
      if key == "module_name":
        # Don't patch `module_name` in case the user implemented it as a
        # computed property.
        continue

      elif key.startswith("__") and key != "__call__":
        # Don't patch methods like `__getattr__` or `__del__`.
        continue

      elif isinstance(value, property):
        # TODO(tomhennigan) Preserve the type of property subclasses.
        clsdict[key] = property(
            value.fget if not value.fget else wrap_method(key, value.fget),
            value.fset if not value.fset else wrap_method(key, value.fset),
            value.fdel if not value.fdel else wrap_method(key, value.fdel),
            doc=value.__doc__)

      elif inspect.isfunction(value):
        # We defer patching methods until after the type is created such that we
        # can trigger the descriptor binding them to the class.
        method_names.append(key)

    clsdict.setdefault(
        "__repr__",
        lambda module: module._auto_repr)  # pylint: disable=protected-access

    cls = super(ModuleMetaclass, mcs).__new__(mcs, name, bases, clsdict)

    for method_name in method_names:
      # Note: the code below is subtle, we need to ensure that we're wrapping
      # the method bound to the class. In some cases (e.g. `wrapt`) this is
      # important since the method can trigger different behavior when it is
      # bound (e.g. in wrapt `FunctionWrapper.__get__(None, cls)` produces a
      # `BoundFunctionWrapper` which in turn populates the `instance` argument
      # to decorator functions using args[0]).
      # Equivalent to: `cls.__dict__[method_name].__get__(None, cls)`
      method = getattr(cls, method_name)
      method = wrap_method(method_name, method)
      setattr(cls, method_name, method)

    return cls

  def __call__(cls: Type[T], *args, **kwargs) -> T:
    # Call new such that we have an un-initialized module instance that we can
    # still reference even if there is an exception during __init__. This is
    # needed such that we can make sure the name_scope constructed in __init__
    # is closed even if there is an exception.

    # NOTE: We disable pytype since (somewhat surprisingly) this method is bound
    # with the new class and not the metaclass.
    module = cls.__new__(cls, *args, **kwargs)  # pytype: disable=wrong-arg-types

    # Now attempt to initialize the object.
    init = wrap_method("__init__", cls.__init__)
    init(module, *args, **kwargs)

    module._auto_repr = utils.auto_repr(cls, *args, **kwargs)  # pylint: disable=protected-access

    ran_super_ctor = hasattr(module, "module_name")
    if not ran_super_ctor:
      raise ValueError(
          "Constructing an hk.Module without calling the super constructor "
          "is not supported. Add the following as the first line in your "
          "__init__ method:\n\nsuper(%s, self).__init__()" % cls.__name__)

    return module


class MethodContext(NamedTuple):
  r"""Read only state showing the calling context for a method.

  For example, let's define two interceptors and print the values in the context.
  Additionally, we will make the first interceptor conditionally short circuit,
  since interceptors stack and are run in order, an earlier interceptor can
  decide to call the next interceptor, or short circuit and call the underlying
  method directly:

  >>> module = hk.Linear(1, name="method_context_example")
  >>> short_circuit = False

  >>> def my_interceptor_1(next_fun, args, kwargs, context):
  ...   print('running my_interceptor_1')
  ...   print('- module.name: ', context.module.name)
  ...   print('- method_name: ', context.method_name)
  ...   if short_circuit:
  ...     return context.orig_method(*args, **kwargs)
  ...   else:
  ...     return next_fun(*args, **kwargs)
  >>> def my_interceptor_2(next_fun, args, kwargs, context):
  ...   print('running my_interceptor_2')
  ...   print('- module.name: ', context.module.name)
  ...   print('- method_name: ', context.method_name)
  ...   return next_fun(*args, **kwargs)

  When ``short_circuit=False`` the two interceptors will run in order:

  >>> with hk.experimental.intercept_methods(my_interceptor_1), \
  ...      hk.experimental.intercept_methods(my_interceptor_2):
  ...   _ = module(jnp.ones([1, 1]))
  running my_interceptor_1
  - module.name:  method_context_example
  - method_name:  __call__
  running my_interceptor_2
  - module.name:  method_context_example
  - method_name:  __call__

  Setting ``short_circuit=True`` will cause the first interceptor to call the
  original method (rather than ``next_fun`` which will trigger the next
  interceptor):

  >>> short_circuit = True
  >>> with hk.experimental.intercept_methods(my_interceptor_1), \
  ...      hk.experimental.intercept_methods(my_interceptor_2):
  ...   _ = module(jnp.ones([1, 1]))
  running my_interceptor_1
  - module.name:  method_context_example
{'fraction_non_alphanumeric': 0.10472559879162771, 'fraction_numerical': 0.0035963461123498523, 'mean_word_length': 3.5025906735751295, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1960197', 'n_tokens_mistral': 4112, 'n_tokens_neox': 4018, 'n_words': 1375}