# Source code for pde.tools.docstrings

"""
Methods for automatic transformation of docstrings

.. autosummary::
:nosignatures:

get_text_block
replace_in_docstring
fill_in_docstring

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""

import re
import textwrap
from typing import Optional, TypeVar

DOCSTRING_REPLACEMENTS = {
# description of function arguments
"ARG_BOUNDARIES_INSTANCE": """
Specifies the boundary conditions applied to the field. This must be an
instance of :class:~pde.grids.boundaries.axes.Boundaries, which can be
created from various data formats using the class method
:func:~pde.grids.boundaries.axes.Boundaries.from_data.
""",
"ARG_BOUNDARIES": """
Boundary conditions are generally given as a list with one condition for each
axis. For periodic axis, only periodic boundary conditions are allowed
(indicated by 'periodic' and 'anti-periodic'). For non-periodic axes, different
boundary conditions can be specified for the lower and upper end (using a tuple
of two conditions). For instance, Dirichlet conditions enforcing a value NUM
(specified by {'value': NUM}) and Neumann conditions enforcing the value DERIV
for the derivative in the normal direction (specified by {'derivative':
DERIV}) are supported. Note that the special value 'natural' imposes periodic
boundary conditions for periodic axis and a vanishing derivative otherwise.
:ref:boundaries documentation <documentation-boundaries>.
""",
"ARG_TRACKER_INTERVAL": """
Determines how often the tracker interrupts the simulation. Simple
numbers are interpreted as durations measured in the simulation time
variable. Alternatively, a string using the format 'hh:mm:ss' can be
used to give durations in real time. Finally, instances of the classes
defined in :mod:~pde.trackers.interrupts can be given for more control.
""",
"ARG_PLOT_QUANTITIES": """
A 2d list of quantities that are shown in a rectangular arrangement.
If quantities is a simple list, the panels will be rendered as a
single row.
Each panel is defined by a dictionary, where the mandatory item 'source'
defines what is being shown. Here, an integer specifies the component
that is extracted from the field while a function is evaluate with the
full state as an input and the result is shown.
Additional items in the dictionary can be 'title' (setting the title of
the panel), 'scale' (defining the color range shown; these are typically
two numbers defining the lower and upper bound, but if only one is given
the range [0, scale] is assumed), and 'cmap' (defining the colormap
being used).
""",
"ARG_PLOT_SCALE": """
Flag determining how the range of the color scale is determined. In the
simplest case a tuple of numbers marks the lower and upper end of the
scalar values that will be shown. If only a single number is supplied,
the range starts at zero and ends at the given number. Additionally, the
special value 'automatic' determines the range from the range of scalar
values.
""",
# descriptions of the discretization and the symmetries
"DESCR_CYLINDRICAL_GRID": r"""
The cylindrical grid assumes polar symmetry, so that fields only depend
on the radial coordinate r and the axial coordinate z. Here, the
first axis is along the radius, while the second axis is along the axis
of the cylinder. The radial discretization is defined as
:math:r_i = (i + \frac12) \Delta r for :math:i=0, \ldots, N_r-1.
""",
"DESCR_POLAR_GRID": r"""
The polar grid assumes polar symmetry, so that fields only depend on the
radial coordinate r. The radial discretization is defined as
:math:r_i = r_\mathrm{min} + (i + \frac12) \Delta r for
:math:i=0, \ldots, N_r-1,  where :math:r_\mathrm{min} is the radius
of the inner boundary, which is zero by default. Note that the radius of
the outer boundary is given by
:math:r_\mathrm{max} = r_\mathrm{min} + N_r \Delta r.
""",
"DESCR_SPHERICAL_GRID": r"""
The spherical grid assumes spherical symmetry, so that fields only
depend on the radial coordinate r. The radial discretization is
defined as :math:r_i = r_\mathrm{min} + (i + \frac12) \Delta r for
:math:i=0, \ldots, N_r-1,  where :math:r_\mathrm{min} is the radius
of the inner boundary, which is zero by default. Note that the radius of
the outer boundary is given by
:math:r_\mathrm{max} = r_\mathrm{min} + N_r \Delta r.
""",
# notes in the docstring
"WARNING_EXEC": r"""
This implementation uses :func:exec and should therefore not be used
in a context where malicious input could occur.
""",
}
DOCSTRING_REPLACEMENTS["ARG_BOUNDARIES_OPTIONAL"] = (
DOCSTRING_REPLACEMENTS["ARG_BOUNDARIES"]
+ "If the special value None is given, no boundary conditions are enforced. The "
"user then needs to ensure that the ghost cells are set accordingly. "
)
DOCSTRING_REPLACEMENTS = {k: v[1:-1] for k, v in DOCSTRING_REPLACEMENTS.items()}

[docs]def get_text_block(identifier: str) -> str:
"""return a single text block

Args:
identifier (str): The name of the text block

Returns:
str: the text block as one long line.
"""
raw_text = DOCSTRING_REPLACEMENTS[identifier]
return "".join(textwrap.dedent(raw_text))

TFunc = TypeVar("TFunc")

[docs]def replace_in_docstring(
f: TFunc, token: str, value: str, docstring: Optional[str] = None
) -> TFunc:
"""replace a text in a docstring using the correct indentation

Args:
f (callable): The function with the docstring to handle
token (str): The token to search for
value (str): The replacement string
docstring (str): A docstring that should be used instead of f.__doc__

Returns:
callable: The function with the modified docstring
"""
# initialize textwrapper for formatting docstring

def repl(matchobj) -> str:
"""helper function replacing token in docstring"""
bare_text = textwrap.dedent(value).strip()
return textwrap.indent(bare_text, matchobj.group(1))

if docstring is None:
docstring = f.__doc__

# replace the token with the correct indentation
f.__doc__ = re.sub(
f"^([ \t]*){token}",
repl,
docstring,  # type: ignore
flags=re.MULTILINE,
)

return f

[docs]def fill_in_docstring(f: TFunc) -> TFunc:
"""decorator that replaces text in the docstring of a function"""
tw = textwrap.TextWrapper(
width=80, expand_tabs=True, replace_whitespace=True, drop_whitespace=True
)

for name, value in DOCSTRING_REPLACEMENTS.items():

def repl(matchobj) -> str:
"""helper function replacing token in docstring"""
tw.initial_indent = tw.subsequent_indent = matchobj.group(1)
return tw.fill(textwrap.dedent(value))

token = "{" + name + "}"
f.__doc__ = re.sub(
f"^([ \t]*){token}",
repl,
f.__doc__,  # type: ignore
flags=re.MULTILINE,
)
return f