proplot-dev · lukelbd · Sep 2, 2019 · Sep 2, 2019 · Sep 2, 2019 · Sep 3, 2019
diff --git a/docs/colors/colors_10_0.svg b/docs/colors/colors_10_0.svg
diff --git a/docs/colors/colors_36_0.svg b/docs/colors/colors_36_0.svg
diff --git a/docs/colors/tutorial_105_0.svg b/docs/colors/tutorial_105_0.svg
diff --git a/docs/colors/tutorial_112_1.svg b/docs/colors/tutorial_112_1.svg
diff --git a/docs/colors/tutorial_114_0.svg b/docs/colors/tutorial_114_0.svg
diff --git a/docs/colors/tutorial_114_1.svg b/docs/colors/tutorial_114_1.svg
diff --git a/docs/colors/tutorial_119_0.svg b/docs/colors/tutorial_119_0.svg
diff --git a/docs/colors/tutorial_119_1.svg b/docs/colors/tutorial_119_1.svg
diff --git a/docs/colors/tutorial_126_0.svg b/docs/colors/tutorial_126_0.svg
diff --git a/docs/colors/tutorial_131_0.svg b/docs/colors/tutorial_131_0.svg
diff --git a/docs/colors/tutorial_133_0.svg b/docs/colors/tutorial_133_0.svg
diff --git a/docs/colors/tutorial_138_0.svg b/docs/colors/tutorial_138_0.svg
diff --git a/docs/colors/tutorial_138_1.svg b/docs/colors/tutorial_138_1.svg
diff --git a/docs/colors/tutorial_157_0.svg b/docs/colors/tutorial_157_0.svg
diff --git a/docs/colors/tutorial_157_1.svg b/docs/colors/tutorial_157_1.svg
diff --git a/docs/colors/tutorial_160_0.svg b/docs/colors/tutorial_160_0.svg
diff --git a/docs/colors/tutorial_163_0.svg b/docs/colors/tutorial_163_0.svg
diff --git a/docs/colors/tutorial_166_0.svg b/docs/colors/tutorial_166_0.svg
diff --git a/docs/colors/tutorial_169_1.svg b/docs/colors/tutorial_169_1.svg
diff --git a/docs/colors/tutorial_172_1.svg b/docs/colors/tutorial_172_1.svg
diff --git a/docs/colors/tutorial_175_0.svg b/docs/colors/tutorial_175_0.svg
diff --git a/docs/colors/tutorial_96_0.svg b/docs/colors/tutorial_96_0.svg
diff --git a/docs/plots/plots_11_0.svg b/docs/plots/plots_11_0.svg
diff --git a/proplot/subplots.py b/proplot/subplots.py
@@ -1,10 +1,9 @@
 #!/usr/bin/env python3
 """
-The starting point for creating custom ProPlot figures and axes.
-The `subplots` function is all you'll need to directly use here.
-It returns a `Figure` instance and an `axes_grid` container of
-`~proplot.axes.Axes` axes, whose positions are controlled by the
-`FlexibleGridSpec` class.
+The `subplots` function is your starting point for creating custom
+ProPlot figures and axes. It returns a `Figure` instance and an
+`axes_grid` container of `~proplot.axes.Axes` axes, whose positions are
+controlled by the `FlexibleGridSpec` class.
 
 .. raw:: html
 
@@ -74,6 +73,7 @@
 import matplotlib.figure as mfigure
 import matplotlib.transforms as mtransforms
 import matplotlib.gridspec as mgridspec
+from numbers import Integral
 try:
     import matplotlib.backends.backend_macosx as mbackend
 except ImportError:
@@ -1457,12 +1457,78 @@ def _update_suptitle(self, title, **kwargs):
         if kwargs:
             self._suptitle.update(kwargs)
 
-    def add_subplot(self, *args, **kwargs):
-        """Issues warning for new users that try to call
-        `~matplotlib.figure.Figure.add_subplot` manually."""
-        if self._locked:
-            warnings.warn('Using "fig.add_subplot()" with ProPlot figures may result in unexpected behavior. Please use "proplot.subplots()" instead.')
-        ax = super().add_subplot(*args, **kwargs)
+    def add_subplot(self, *args,
+        proj=None, projection=None, basemap=False,
+        **kwargs):
+        """
+        Adds subplot using the existing figure gridspec.
+
+        Parameters
+        ----------
+        *args
+            If `~matplotlib.gridspec.SubplotSpec` instance, must be a child
+            of the main subplot gridspec!
+
+            If a 3-digit integer or a tuple indicating (nrows, ncols, index),
+            the geometry must match the geometry of the input gridspec!
+        proj, projection : str, optional
+            The registered matplotlib projection name, or a basemap or cartopy
+            map projection name. For valid map projection names, see the
+            :ref:`Table of projections`.
+        """
+        # TODO: Consider permitting add_subplot?
+        # Copied from matplotlib add_subplot
+        if not len(args):
+            args = (1, 1, 1)
+        if len(args) == 1 and isinstance(args[0], Integral):
+            if not 100 <= args[0] <= 999:
+                raise ValueError(f'Integer subplot specification must be a three-digit number, not {args[0]!r}.')
+            args = tuple(map(int, str(args[0])))
+        # Copied from SubplotBase __init__ and modified to enforce restrictions
+        gridspec = self._gridspec_main
+        subplotspec = None
+        if len(args) == 1:
+            if isinstance(args[0], mgridspec.SubplotSpec):
+                subplotspec = args[0]
+            else:
+                try:
+                    s = str(int(args[0]))
+                    rows, cols, num = map(int, s)
+                except ValueError:
+                    raise ValueError(f'Single argument to subplot must be a 3-digit integer, not {args[0]!r}.')
+        elif len(args) == 3:
+            rows, cols, num = args
+        else:
+            raise ValueError(f'Illegal argument(s) to add_subplot: {args}')
+        if subplotspec is None:
+            rows = int(rows)
+            cols = int(cols)
+            if isinstance(num, tuple) and len(num) == 2:
+                num = [int(n) for n in num]
+            else:
+                if num < 1 or num > rows*cols:
+                    raise ValueError(f'num must be 1 <= num <= {rows*cols}, not {num}')
+        if (rows, cols) != gridspec.get_active_geometry():
+            raise ValueError(f'Input arguments {args!r} conflict with existing gridspec geometry of {rows} rows, {cols} columns.')
+        if not isinstance(num, tuple):
+            num = (num, num)
+        subplotspec = gridspec[(num[0] - 1):num[1]]
+
+        # The default is CartesianAxes
+        proj = _notNone(proj, projection, 'cartesian', names=('proj','projection'))
+        # Builtin matplotlib polar axes, just use my overridden version
+        if proj == 'polar':
+            proj = 'polar2'
+        # Custom Basemap and Cartopy axes
+        # TODO: Have Proj return all unused keyword args, with a
+        # map_projection = obj entry, and maybe hide the Proj constructor as
+        # an argument processing utility?
+        elif proj != 'cartesian':
+            kwargs = projs.Proj(proj, basemap=basemap, **kwargs)
+            proj = 'basemap' if basemap else 'cartopy'
+
+        # Initialize
+        ax = super().add_subplot(subplotspec, projection=proj, **kwargs)
         return ax
 
     def colorbar(self, *args,
@@ -1844,7 +1910,9 @@ def subplots(array=None, ncols=1, nrows=1,
         for the *x* axis, *y* axis, or both axes. Only has an effect when
         `spanx`, `spany`, or `span` are ``False``.
     proj, projection : str or dict-like, optional
-        The map projection name. The argument is interpreted as follows.
+        The registered matplotlib projection name, or a basemap or cartopy
+        map projection name. For valid map projection names, see the
+        :ref:`Table of projections`. Argument is interpreted as follows.
 
         * If string, this projection is used for all subplots. For valid
           names, see the :ref:`Table of projections`.