old_contents: """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
  
new_contents: """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
  
subject: Add decorator for saving plts