Use enum etc

README.md

@@ -34,14 +34,13 @@ advocate for the latter ([proof](https://agustinus.kristia.de/blog/plotting/)).
 <br />
 
 ```diff
-import pub_ready_plots
+import pub_ready_plots as prp
 
 ...
 
-pub_ready_plots.get_context(
-    ...
--   layout="icml",
-+   layout="poster-landscape",
+prp.get_context(
+-   layout=prp.Layout.ICML,
++   layout=prp.Layout.POSTER_LANDSCAPE,
     ...
 )
 
@@ -59,18 +58,12 @@ pip install pub-ready-plots
 ## Quick usage
 
 ```python
-import pub_ready_plots
+import pub_ready_plots as prp
 
-with pub_ready_plots.get_context(
-    width_frac=1,  # Multiplier for `\linewidth`
-    height_frac=0.15,  # Multiplier for `\textheight`
-    layout="icml",  # or "iclr", "neurips", "aistats", "uai", "tmlr", "poster-portrait", "poster-landscape"
-    single_col=False,  # only works for the "icml", "aistats", "uai" layouts
-    nrows=1,  # depending on your subplots, default = 1
-    ncols=2,  # depending on your subplots, default = 1
-    override_rc_params={"lines.linewidth": 4.123},  # Overriding rcParams
-    sharey=True,  # Additional keyword args for `plt.subplots`
-) as (fig, axs):
+# Wrap you current plotting script with this `with` statement.
+# By default, this will create a full-width, 0.15*\textheight plot that conforms
+# to the ICLR template.
+with prp.get_context(layout=prp.Layout.ICLR) as (fig, axs):
     # Do whatever you want with `fig` and `axs`
     ...
 
@@ -91,6 +84,7 @@ Then in your LaTeX file, include the plot as follows:
 
 That's it! But you should use TikZ more.
 Anyway, see the full, runnable example in [`examples/simple_plot.py`](https://github.com/wiseodd/pub-ready-plots/blob/master/examples/simple_plot.py)
+See [here](#all-available-options) for available options for `get_context()`!
 
 > [!TIP]
 > I recommend using this library in conjunction with
@@ -100,17 +94,37 @@ Anyway, see the full, runnable example in [`examples/simple_plot.py`](https://gi
 
 ## Advanced usages
 
+### All available options
+
+```python
+import pub_ready_plots as prp
+
+with prp.get_context(
+    layout=prp.Layout.ICML,  # check `Layout` for all available layouts
+    width_frac=1,  # multiplier for `\linewidth`
+    height_frac=0.15,  # multiplier for `\textheight`
+    single_col=False,  # only works for the "icml", "aistats", "uai" layouts
+    nrows=1,  # depending on your subplots, default = 1
+    ncols=2,  # depending on your subplots, default = 1
+    override_rc_params={"lines.linewidth": 4.123},  # Overriding rcParams
+    sharey=True,  # Additional keyword args for `plt.subplots`
+) as (fig, axs):
+    ...
+
+    fig.savefig("filename.pdf")
+```
+
 ### Creating plots for `\wrapfigure`
 
 Say we want to have an inline figure of size `0.4\textwidth` and
 height `0.15\textheight` in our NeurIPS paper.
 Then all we have to do is the following:
 
 ```python
-import pub_ready_plots
+import pub_ready_plots as prp
 
-with pub_ready_plots.get_context(
-    width_frac=0.4, height_frac=0.15, layout="neurips",
+with prp.get_context(
+    layout=prp.Layout.NEURIPS, width_frac=0.4, height_frac=0.15,
 ) as (fig, axs):
     # Your plot here!
     ...

examples/advanced_usage.py

@@ -2,15 +2,15 @@
 import numpy as np
 from matplotlib.axes import Axes
 
-import pub_ready_plots
+import pub_ready_plots as prp
 
 ########################################################################################
 # User-specified rcParams override
 ########################################################################################
-with pub_ready_plots.get_context(
+with prp.get_context(
+    layout=prp.Layout.ICLR,
     width_frac=1,
     height_frac=0.15,
-    layout="iclr",
     override_rc_params={"lines.linewidth": 5},  # Pass your style overrides here!
 ) as (fig, ax):
     assert isinstance(ax, Axes)
@@ -28,10 +28,10 @@
 ########################################################################################
 # Manual, most-flexible way to use this library
 ########################################################################################
-rc_params, fig_width_in, fig_height_in = pub_ready_plots.get_mpl_rcParams(
+rc_params, fig_width_in, fig_height_in = prp.get_mpl_rcParams(
+    layout=prp.Layout.POSTER_PORTRAIT,
     width_frac=1,
     height_frac=0.15,
-    layout="poster-portrait",
     single_col=False,
 )
 
@@ -47,6 +47,8 @@
 
 x = np.linspace(-1, 1, 100)
 
+assert isinstance(axs, np.ndarray)
+
 axs[0].plot(x, np.sin(x))
 axs[0].set_title("Sine")
 axs[0].set_xlabel(r"$x$")

examples/simple_plot.py

@@ -1,16 +1,18 @@
 import numpy as np
 from matplotlib.axes import Axes
 
-import pub_ready_plots
+import pub_ready_plots as prp
 
 ########################################################################################
 # Single plot (i.e. no subplots)
 ########################################################################################
-with pub_ready_plots.get_context(
-    width_frac=1,  # between 0 and 1
-    height_frac=0.15,  # between 0 and 1
-    layout="iclr",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
-) as (fig, ax):
+with (
+    prp.get_context(
+        layout=prp.Layout.ICLR,  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
+        width_frac=1,  # between 0 and 1
+        height_frac=0.15,  # between 0 and 1
+    ) as (fig, ax)
+):
     # Just like in `plt.subplots`, `ax` is a matplotlib Axes if
     # nrows & ncols are not specified (both default to 1).
     assert isinstance(ax, Axes)
@@ -27,15 +29,17 @@
 ########################################################################################
 # Multiple subplots
 ########################################################################################
-with pub_ready_plots.get_context(
-    width_frac=1,  # between 0 and 1
-    height_frac=0.15,  # between 0 and 1
-    nrows=1,  # depending on your subplots
-    ncols=2,  # depending on your subplots
-    layout="iclr",  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
-    single_col=False,  # only works for the "icml" layout
-    sharey=True,  # Additional keyword args for `plt.subplots`
-) as (fig, axs):
+with (
+    prp.get_context(
+        layout=prp.Layout.ICLR,  # or "iclr", "neurips", "poster-portrait", "poster-landscape"
+        width_frac=1,  # between 0 and 1
+        height_frac=0.15,  # between 0 and 1
+        nrows=1,  # depending on your subplots
+        ncols=2,  # depending on your subplots
+        single_col=False,  # only works for the "icml" layout
+        sharey=True,  # Additional keyword args for `plt.subplots`
+    ) as (fig, axs)
+):
     # If `nrows` or `ncols` are not 1, `axs` is a NumPy array containing Axes'
     assert isinstance(axs, np.ndarray)
 

makefile

@@ -1,3 +1,6 @@
+test:
+	uv run pytest --cov
+
 ruff:
 	-uv run ruff format
 	@uv run ruff check --fix

pub_ready_plots/__init__.py

@@ -1,3 +1,4 @@
 from pub_ready_plots.pub_ready_plots import get_context, get_mpl_rcParams
+from pub_ready_plots.styles import Layout
 
-__all__ = ["get_mpl_rcParams", "get_context"]
+__all__ = ["get_mpl_rcParams", "get_context", "Layout"]

pub_ready_plots/pub_ready_plots.py

@@ -6,22 +6,25 @@
 from matplotlib.figure import Figure
 from numpy import ndarray
 
-from .styles import PAPER_FORMATS, Style
+from .styles import PAPER_FORMATS, Layout, Style
 
 
 @contextmanager
 def get_context(
+    layout: Layout,
     width_frac: float = 1,
     height_frac: float = 0.15,
-    layout: str = "neurips",
     single_col: bool = False,
     nrows: int = 1,
     ncols: int = 1,
     override_rc_params: dict[str, Any] = dict(),
     **kwargs: Any,
 ) -> Generator[tuple[Figure, Union[Axes, ndarray[Any, Any]]], None, None]:
     rc_params, fig_width_in, fig_height_in = get_mpl_rcParams(
-        width_frac, height_frac, layout, single_col
+        layout=layout,
+        width_frac=width_frac,
+        height_frac=height_frac,
+        single_col=single_col,
     )
     rc_params.update(override_rc_params)
 
@@ -32,9 +35,9 @@ def get_context(
 
 
 def get_mpl_rcParams(
+    layout: Layout,
     width_frac: float = 1,
     height_frac: float = 0.15,
-    layout: str = "neurips",
     single_col: bool = False,
 ) -> tuple[dict[str, Any], float, float]:
     """Get matplotlib rcParams dict and fig width & height in inches, depending on the
@@ -44,7 +47,7 @@ def get_mpl_rcParams(
 
     ```python
         rc_params, fig_width_in, fig_height_in = pub_ready_plots.get_mpl_rcParams(
-            width_frac=fig_width_frac, height_frac=fig_height_frac, layout="icml"
+            layout=Layout.ICML, width_frac=fig_width_frac, height_frac=fig_height_frac
         )
         plt.rcParams.update(rc_params)
 
@@ -69,12 +72,13 @@ def get_mpl_rcParams(
     The arg. `width=\\linewidth` is important!
 
     Args:
+        layout: The LaTeX template used. Possible values are Layout.ICML, Layout.NeurIPS,
+            Layout.ICLR, Layout.AISTATS, Layout.UAI, Layout.JMLR, Layout.TMLR,
+            Layout.POSTER_PORTRAIT (A1, 2-column), and Layout.POSTER_LANDSCAPE (A0, 3-col).
         width_frac: Fraction of `\\linewidth` as the figure width. Usually set to 1.
         height_frac: Fraction of `\\textheight` as the figure height. Try 0.175.
-        layout: The LaTeX template used. Possible values are "icml", "iclr", "neurips",
-            "jmlr", "poster-portrait" (A1, 2-column), and "poster-landscape" (A0, 3-col).
-        single_col: Whether the plot is single column in a layout that has two columns
-            (e.g. ICML). Not supported for any other layout.
+        single_col: Whether the plot is single column in a layout that has multiple columns
+            (e.g. ICML, posters). Not supported for any other layout.
 
     Returns:
         rc_params: Matplotlib key-value rc-params. Use it via
@@ -86,14 +90,23 @@ def get_mpl_rcParams(
     if (width_frac <= 0 or width_frac > 1) or (height_frac <= 0 or height_frac > 1):
         raise ValueError("Both `width_frac` and `height_frac` must be between 0 and 1.")
 
-    if layout not in PAPER_FORMATS.keys():
-        raise ValueError(f"Layout must be in {list(PAPER_FORMATS.keys())}.")
-
-    if layout not in ["icml", "aistats", "uai"] and single_col:
-        raise ValueError("Double-column is only supported for ICML, AISTATS, and UAI.")
+    if (
+        layout
+        not in [
+            Layout.ICML,
+            Layout.AISTATS,
+            Layout.UAI,
+            Layout.POSTER_PORTRAIT,
+            Layout.POSTER_LANDSCAPE,
+        ]
+        and single_col
+    ):
+        raise ValueError(
+            "Double-column is only supported for ICML, AISTATS, UAI, POSTER_PORTRAIT, and POSTER_LANDSCAPE."
+        )
 
     format: Style = PAPER_FORMATS[layout]
-    is_poster = "poster" in layout
+    is_poster: bool = "poster" in layout._name_.lower()
 
     rc_params = {
         "text.usetex": False,