densitty 0.8.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

densitty/detect.py

@@ -9,14 +9,14 @@ from types import MappingProxyType
 from typing import Any, Callable, Optional, Sequence
 import time
 
-from . import ansi, ascii_art, binning, lineart, truecolor
-from . import plot as plotmodule
-from .util import FloatLike, ValueRange
+from . import ansi, ascii_art, axis, binning, colorbar, lineart, plotting, smoothing, truecolor
+from .util import FloatLike, ValueRange, make_value_range
 
 if sys.platform == "win32":
     # pylint: disable=import-error
     import ctypes
-    from ctypes.windll.kernel32 import GetConsoleMode, GetStdHandle, SetConsoleMode
+
+    kernel32 = ctypes.windll.kernel32
 else:
     # All other platforms should have TERMIOS available
     import fcntl
@@ -86,10 +86,10 @@ if sys.platform == "win32":
         """Windows-based wrapper to avoid control code output to stdout"""
         prev_stdin_mode = ctypes.wintypes.DWORD(0)
         prev_stdout_mode = ctypes.wintypes.DWORD(0)
-        GetConsoleMode(GetStdHandle(-10), ctypes.byref(prev_stdin_mode))
-        SetConsoleMode(GetStdHandle(-10), 0)
-        GetConsoleMode(GetStdHandle(-11), ctypes.byref(prev_stdout_mode))
-        SetConsoleMode(GetStdHandle(-11), 7)
+        kernel32.GetConsoleMode(kernel32.GetStdHandle(-10), ctypes.byref(prev_stdin_mode))
+        kernel32.SetConsoleMode(kernel32.GetStdHandle(-10), 0)
+        kernel32.GetConsoleMode(kernel32.GetStdHandle(-11), ctypes.byref(prev_stdout_mode))
+        kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
 
         # On Windows, don't try to be non-blocking, just read until terminator
         if length is None:
@@ -106,8 +106,8 @@ if sys.platform == "win32":
                     break
             return response
         finally:
-            SetConsoleMode(GetStdHandle(-10), ctypes.byref(prev_stdin_mode))
-            SetConsoleMode(GetStdHandle(-11), ctypes.byref(prev_stdout_mode))
+            kernel32.SetConsoleMode(kernel32.GetStdHandle(-10), ctypes.byref(prev_stdin_mode))
+            kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), ctypes.byref(prev_stdout_mode))
 
 else:
     # Not Windows, so use termios/fcntl:
@@ -304,7 +304,8 @@ def color_support(interactive=True, debug=False) -> ColorSupport:
         return ColorSupport.NONE
 
     env_term = os.environ.get("TERM", "")
-    print(f"$TERM is {env_term}")
+    if debug:
+        print(f"$TERM is {env_term}")
 
     truecolor_terminals = ("truecolor", "xterm-kitty", "xterm-ghostty", "wezterm")
     if env_term in truecolor_terminals:
@@ -362,6 +363,13 @@ def color_support(interactive=True, debug=False) -> ColorSupport:
         return min(ColorSupport.ANSI_24BIT, color_cap)
 
     if curses:
+        # Curses is installed, but it may or may not have been set up. Try and see
+        try:
+            curses.tigetflag("RGB")
+            # If this gets an error, it can be an internal type not derived from Exception
+            # so just catch everything:
+        except:  # pylint: disable=bare-except
+            curses.setupterm()
         if curses.tigetflag("RGB") == 1:
             # ncurses 6.0+ / terminfo added an "RGB" capability to indicate truecolor support
             return min(ColorSupport.ANSI_24BIT, color_cap)
@@ -371,6 +379,8 @@ def color_support(interactive=True, debug=False) -> ColorSupport:
         if curses.tigetnum("colors") == 256:
             return min(ColorSupport.ANSI_8BIT, color_cap)
 
+        curses_colors = curses.tigetnum("colors")  # for use below
+
     if env_term.endswith("-256color") or env_term.endswith("-256"):
         if debug:
             print("Color detect: $TERM suffix in 8b list")
@@ -391,8 +401,9 @@ def color_support(interactive=True, debug=False) -> ColorSupport:
             print("Color detect: using terminal's Device Attributes")
         return da1_color_support(debug)
 
-    if curses and curses.tigetnum("colors") >= 16:
-        return min(ColorSupport.ANSI_4BIT, color_cap)
+    if curses:
+        if curses_colors >= 16:
+            return min(ColorSupport.ANSI_4BIT, color_cap)
 
     return ColorSupport.NONE
 
@@ -415,6 +426,24 @@ FADE_IN = MappingProxyType(
     }
 )
 
+RAINBOW = MappingProxyType(
+    {
+        ColorSupport.NONE: ascii_art.EXTENDED,
+        ColorSupport.ANSI_4BIT: ansi.RAINBOW_16,
+        ColorSupport.ANSI_8BIT: ansi.RAINBOW,
+        ColorSupport.ANSI_24BIT: truecolor.RAINBOW,
+    }
+)
+
+REV_RAINBOW = MappingProxyType(
+    {
+        ColorSupport.NONE: ascii_art.EXTENDED,
+        ColorSupport.ANSI_4BIT: ansi.REV_RAINBOW_16,
+        ColorSupport.ANSI_8BIT: ansi.REV_RAINBOW,
+        ColorSupport.ANSI_24BIT: truecolor.REV_RAINBOW,
+    }
+)
+
 
 def pick_colormap(maps: dict[ColorSupport, Callable]) -> Callable:
     """Detect color support and pick the best color map"""
@@ -422,21 +451,33 @@ def pick_colormap(maps: dict[ColorSupport, Callable]) -> Callable:
     return maps[support]
 
 
-def plot(data, colors=FADE_IN, **plotargs):
-    """Wrapper for plot.Plot that picks colormap from dict"""
+def plot(data, colors=FADE_IN, colorscale=False, **plotargs):
+    """Wrapper for plot.Plot that picks colormap from dict
+
+    Parameters
+     ----------
+     data : Sequence[Sequence[float]]
+            The data to be plotted.
+     colors : Dict mapping color support to color map
+            The color scale (e.g. GRAYSCALE, FADE_IN)
+     colorscale: bool
+            Display a vertical color scale on the right
+            (for horizontal, just directly create a colorbar
+              and output it after outputting the plot)
+     plotargs: any Plot() keyword arguments
+    """
     colormap = pick_colormap(colors)
-    return plotmodule.Plot(data, colormap, **plotargs)
+    the_plot = plotting.Plot(data, colormap, **plotargs)
+    if colorscale:
+        colorbar.add_colorbar(the_plot)
+    return the_plot
 
 
 def histplot2d(
     points: Sequence[tuple[FloatLike, FloatLike]],
-    bins: (
-        int
-        | tuple[int, int]
-        | Sequence[FloatLike]
-        | tuple[Sequence[FloatLike], Sequence[FloatLike]]
-    ) = 10,
+    bins: binning.FullBinsArg = None,
     ranges: Optional[tuple[Optional[ValueRange], Optional[ValueRange]]] = None,
+    bin_size: Optional[FloatLike | tuple[FloatLike, FloatLike]] = None,
     align=True,
     drop_outside=True,
     colors=FADE_IN,
@@ -449,8 +490,9 @@ def histplot2d(
     """Wrapper for binning.histogram2d / plot.Plot to simplify 2-D histogram plotting"""
     binned_data, x_axis, y_axis = binning.histogram2d(
         points,
-        bins,
+        bins=bins,
         ranges=ranges,
+        bin_size=bin_size,
         align=align,
         drop_outside=drop_outside,
         border_line=border_line,
@@ -460,6 +502,82 @@ def histplot2d(
     if scale is True:
         p.upscale()
     elif scale:
-        p.upscale(max_expansion=(scale, scale))
+        p.upscale(max_expansion=scale)
+
+    return p
+
+
+def densityplot2d(
+    points: Sequence[tuple[FloatLike, FloatLike]],
+    kernel: Optional[smoothing.SmoothingFunc] = None,
+    bins: binning.FullBinsArg = None,
+    ranges: Optional[tuple[Optional[ValueRange], Optional[ValueRange]]] = None,
+    align=True,
+    colors=FADE_IN,
+    border_line=True,
+    fractional_tick_pos=False,
+    **plotargs,
+    # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals
+):
+    """Wrapper for smoothing.smooth2d / plot.Plot to simplify 2-D density plots"""
+
+    if not bins:
+        try:
+            terminal_size: Optional[os.terminal_size] = os.get_terminal_size()
+        except OSError:
+            terminal_size = plotting.default_terminal_size
+        if terminal_size is None:
+            raise OSError("No terminal size from os.get_terminal_size()")
+        size_x = terminal_size.columns - 10
+        size_y = terminal_size.lines - 4
+        bins = (size_x, size_y)
+
+    if kernel is None:
+        _, num_bins, bin_centers = binning.expand_bins_arg(bins)
+        if bin_centers:
+            # we were given a list of bin centers, so generate bounds from that:
+            x_bin_range = make_value_range((bin_centers[0][0], bin_centers[0][-1]))
+            y_bin_range = make_value_range((bin_centers[1][0], bin_centers[1][-1]))
+            x_width, y_width = smoothing.pick_kernel_bandwidth(
+                points, bins=num_bins, ranges=(x_bin_range, y_bin_range)
+            )
+        else:
+            x_width, y_width = smoothing.pick_kernel_bandwidth(points, bins=num_bins)
+        kernel = smoothing.gaussian_with_sigmas(x_width, y_width)
+
+    smoothed, x_axis, y_axis = smoothing.smooth2d(
+        points=points,
+        kernel=kernel,
+        bins=bins,
+        ranges=ranges,
+        align=align,
+        border_line=border_line,
+        fractional_tick_pos=fractional_tick_pos,
+    )
+    p = plot(smoothed, colors, x_axis=x_axis, y_axis=y_axis, **plotargs)
 
     return p
+
+
+def grid_heatmap(data, x_labels, y_labels, colors=REV_RAINBOW, max_cell_size=0, **plotargs):
+    """Create a grid-style heatmap, with explicit X and Y labels for each bin"""
+    if max_cell_size == 0:
+        max_x_label = max(len(x) for x in x_labels)
+        max_cell_size = max_x_label + 2
+
+    num_rows = len(data)
+    num_cols = len(data[0])
+    if y_labels and len(y_labels) != num_rows:
+        raise ValueError("Number of Y labels does not match number of rows")
+    if x_labels and len(x_labels) != num_cols:
+        raise ValueError("Number of X labels does not match number of columns")
+
+    # value range along axis is arbitrary, but is needed to locate the labels along the axis
+    # so just use integers 0..N-1
+    x_axis = axis.Axis((0, num_cols - 1), dict(enumerate(x_labels)))
+    y_axis = axis.Axis((0, num_rows - 1), dict(enumerate(y_labels)))
+
+    plt = plot(data, colors=colors, x_axis=x_axis, y_axis=y_axis, flip_y=False, **plotargs)
+    plt.upscale(max_expansion=max_cell_size)
+
+    return plt

densitty/{plot.py → plotting.py}

@@ -1,13 +1,20 @@
 """Two-dimensional histogram (density plot) with textual output."""
 
 import dataclasses
-from itertools import chain, zip_longest
 import os
 import sys
-from typing import Callable, Optional, Sequence
+import typing
+from itertools import chain, zip_longest
+from typing import Any, Callable, Optional, Sequence
+
+from . import ansi, lineart
 
-from . import ansi, axis, lineart
-from .util import FloatLike
+if typing.TYPE_CHECKING:
+    from .axis import Axis
+    from .util import FloatLike
+else:
+    Axis = Any
+    FloatLike = Any
 
 # pylint: disable=invalid-name
 # User can set this to provide a default if os.terminal_size() fails:
@@ -27,9 +34,17 @@ class Plot:
     font_mapping: dict = dataclasses.field(default_factory=lambda: lineart.basic_font)
     min_data: Optional[FloatLike] = None
     max_data: Optional[FloatLike] = None
-    x_axis: Optional[axis.Axis] = None
-    y_axis: Optional[axis.Axis] = None
+    x_axis: Optional[Axis] = None
+    y_axis: Optional[Axis] = None
     flip_y: bool = True  # put the first row of data at the bottom of the output
+    to_right: Optional["Plot" | Sequence[str]] = None
+    right_padding: Optional[str] = None
+
+    def data_limits(self):
+        """Return (min,max) of the plot data."""
+        min_data = min(min(line) for line in self.data) if self.min_data is None else self.min_data
+        max_data = max(max(line) for line in self.data) if self.max_data is None else self.max_data
+        return (min_data, max_data)
 
     def as_ascii(self):
         """Output using direct characters (ASCII-art)."""
@@ -66,8 +81,7 @@ class Plot:
         Also flips data if requested
         """
 
-        min_data = min(min(line) for line in self.data) if self.min_data is None else self.min_data
-        max_data = max(max(line) for line in self.data) if self.max_data is None else self.max_data
+        min_data, max_data = self.data_limits()
         data_scale = max_data - min_data
         if data_scale == 0:
             # all data has the same value (or we were told it does)
@@ -84,6 +98,17 @@ class Plot:
         """Are there two pixels per output character?"""
         return self.is_color() and self.render_halfheight
 
+    def left_margin(self):
+        """Returns the column of the first plotted data"""
+        if self.y_axis:
+            if self.is_halfheight():
+                num_rows = (len(self.data) + 1) // 2
+            else:
+                num_rows = len(self.data)
+            y_axis_lines = self.y_axis.render_as_y(num_rows, False, False, False)
+            return lineart.display_len(y_axis_lines[0])
+        return 0
+
     def as_strings(self):
         """Scale data to 0..1 range and feed it through the appropriate output function"""
         if self.is_halfheight():
@@ -98,22 +123,34 @@ class Plot:
 
         if self.y_axis:
             axis_lines = self.y_axis.render_as_y(num_rows, False, bool(self.x_axis), self.flip_y)
+            left_margin = lineart.display_len(axis_lines[0])
         else:
             axis_lines = ["" for _ in range(num_rows + bool(self.x_axis))]
+            left_margin = 0
 
-        left_margin = lineart.display_len(axis_lines[0])
         if self.x_axis:
             x_ticks, x_labels = self.x_axis.render_as_x(num_cols, left_margin)
             axis_lines[-1] = lineart.merge_lines(x_ticks, axis_lines[-1])
             axis_lines += [x_labels]
 
-        for frame, plot_line in zip_longest(axis_lines, plot_lines, fillvalue=""):
-            yield frame.translate(self.font_mapping) + plot_line
+        if isinstance(self.to_right, Plot):
+            lines_to_right = self.to_right.as_strings()
+        elif self.to_right is not None:
+            lines_to_right = self.to_right
+        else:
+            lines_to_right = []
+
+        padding = self.right_padding if self.right_padding is not None else ""
+
+        for frame, plot_line, glued_on in zip_longest(
+            axis_lines, plot_lines, lines_to_right, fillvalue=""
+        ):
+            yield frame.translate(self.font_mapping) + plot_line + padding + glued_on
 
-    def show(self, printer=print):
+    def show(self, prefix="", printer=print):
         """Simple helper function to output/print a plot"""
         for line in self.as_strings():
-            printer(line)
+            printer(prefix + line)
 
     def _compute_scaling_multipliers(
         self,
@@ -156,24 +193,29 @@ class Plot:
 
     def upscale(
         self,
-        max_size: tuple[int, int] = (0, 0),
-        max_expansion: tuple[int, int] = (3, 3),
-        keep_aspect_ratio: bool = False,
+        max_size: int | tuple[int, int] = 0,
+        max_expansion: int | tuple[int, int] = 3,
+        keep_aspect_ratio: bool = True,
     ):
         """Scale up 'data' by repeating lines and values within lines.
 
         Parameters
         ----------
-        max_size : tuple (int, int)
+        max_size : int or tuple (int, int)
                    If positive: Maximum number of columns, maximum number of rows
                    If zero: Use terminal size
                    If negative: Use as offset from terminal size
                    Default: Based on terminal size (0).
-        max_expansion : tuple (int, int)
-                   maximum expansion factor in each direction. Default (3,3). 0=> No maximum
+        max_expansion : int or tuple (int, int)
+                   maximum expansion factor in each direction. Default (3,3). None=> No maximum
         keep_aspect_ratio : bool
                    Require that X and Y scaling are equal.
         """
+        if isinstance(max_size, int):
+            max_size = (max_size, max_size)
+
+        if isinstance(max_expansion, int) or max_expansion is None:
+            max_expansion = (max_expansion, max_expansion)
 
         float_mult = self._compute_scaling_multipliers(max_size, keep_aspect_ratio)
         col_mult = max(int(float_mult[0]), 1)
@@ -198,4 +240,20 @@ class Plot:
 
         # repeat each of those by the row multiplier
         self.data = repeat_each(x_expanded, row_mult)
+
+        # If we have axes, adjust for the new pixel/bin size
+        if self.x_axis:
+            self.x_axis.upscale(len(self.data[0]), col_mult)
+
+        if self.y_axis:
+            self.y_axis.upscale(len(self.data), row_mult)
+
+        # if we have a glued-on plot to the right (likely a colorbar), scale it up in Y to match
+        if isinstance(self.to_right, Plot):
+            self.to_right.upscale(max_size=0, max_expansion=(1, col_mult), keep_aspect_ratio=False)
         return self
+
+    def glue_on(self, to_right, padding="  "):
+        """Add lines or another plot to the right of this one."""
+        self.right_padding = padding
+        self.to_right = to_right