"""Utility functions for NeuroDSP plots."""

import matplotlib.pyplot as plt

###################################################################################################
###################################################################################################

def check_ax(ax, figsize=None):
    """Check whether a figure axes object is defined, define if not.

    Parameters
    ----------
    ax : matplotlib.Axes or None
        Axes object to check if is defined.

    Returns
    -------
    ax : matplotlib.Axes
        Figure axes object to use.
    """

    if not ax:
        _, ax = plt.subplots(figsize=figsize)

    return ax

"""Utility functions for NeuroDSP plots."""

from functools import wraps
from os.path import join as pjoin

import matplotlib.pyplot as plt

###################################################################################################
###################################################################################################

def check_ax(ax, figsize=None):
    """Check whether a figure axes object is defined, define if not.

    Parameters
    ----------
    ax : matplotlib.Axes or None
        Axes object to check if is defined.

    Returns
    -------
    ax : matplotlib.Axes
        Figure axes object to use.
    """

    if not ax:
        _, ax = plt.subplots(figsize=figsize)

    return ax


def savefig(func):

    @wraps(func)
    def decorated(*args, **kwargs):

        save_fig = kwargs.pop('save_fig', False)
        file_name = kwargs.pop('file_name', None)
        file_path = kwargs.pop('file_path', None)

        func(*args, **kwargs)

        if save_fig:
            full_path = pjoin(file_path, file_name) if file_path else file_name
            plt.savefig(full_path)

    return decorated

Add decorator for saving plts
Add decorator for saving plts

  """Utility functions for NeuroDSP plots."""
+ 
+ from functools import wraps
+ from os.path import join as pjoin
  
  import matplotlib.pyplot as plt
  
  ###################################################################################################
  ###################################################################################################
  
  def check_ax(ax, figsize=None):
      """Check whether a figure axes object is defined, define if not.
  
      Parameters
      ----------
      ax : matplotlib.Axes or None
          Axes object to check if is defined.
  
      Returns
      -------
      ax : matplotlib.Axes
          Figure axes object to use.
      """
  
      if not ax:
          _, ax = plt.subplots(figsize=figsize)
  
      return ax
  
+ 
+ def savefig(func):
+ 
+     @wraps(func)
+     def decorated(*args, **kwargs):
+ 
+         save_fig = kwargs.pop('save_fig', False)
+         file_name = kwargs.pop('file_name', None)
+         file_path = kwargs.pop('file_path', None)
+ 
+         func(*args, **kwargs)
+ 
+         if save_fig:
+             full_path = pjoin(file_path, file_name) if file_path else file_name
+             plt.savefig(full_path)
+ 
+     return decorated
+ 
Add decorator for saving plts
## Code Before:
"""Utility functions for NeuroDSP plots."""

import matplotlib.pyplot as plt

###################################################################################################
###################################################################################################

def check_ax(ax, figsize=None):
    """Check whether a figure axes object is defined, define if not.

    Parameters
    ----------
    ax : matplotlib.Axes or None
        Axes object to check if is defined.

    Returns
    -------
    ax : matplotlib.Axes
        Figure axes object to use.
    """

    if not ax:
        _, ax = plt.subplots(figsize=figsize)

    return ax

## Instruction:
Add decorator for saving plts
## Code After:
"""Utility functions for NeuroDSP plots."""

from functools import wraps
from os.path import join as pjoin

import matplotlib.pyplot as plt

###################################################################################################
###################################################################################################

def check_ax(ax, figsize=None):
    """Check whether a figure axes object is defined, define if not.

    Parameters
    ----------
    ax : matplotlib.Axes or None
        Axes object to check if is defined.

    Returns
    -------
    ax : matplotlib.Axes
        Figure axes object to use.
    """

    if not ax:
        _, ax = plt.subplots(figsize=figsize)

    return ax


def savefig(func):

    @wraps(func)
    def decorated(*args, **kwargs):

        save_fig = kwargs.pop('save_fig', False)
        file_name = kwargs.pop('file_name', None)
        file_path = kwargs.pop('file_path', None)

        func(*args, **kwargs)

        if save_fig:
            full_path = pjoin(file_path, file_name) if file_path else file_name
            plt.savefig(full_path)

    return decorated
