Aluode/PerceptionLabPortable

1import matplotlib.axes as maxes2from matplotlib.artist import Artist3from matplotlib.axis import XAxis, YAxis4 5 6class SimpleChainedObjects:7    def __init__(self, objects):8        self._objects = objects9 10    def __getattr__(self, k):11        _a = SimpleChainedObjects([getattr(a, k) for a in self._objects])12        return _a13 14    def __call__(self, *args, **kwargs):15        for m in self._objects:16            m(*args, **kwargs)17 18 19class Axes(maxes.Axes):20 21    class AxisDict(dict):22        def __init__(self, axes):23            self.axes = axes24            super().__init__()25 26        def __getitem__(self, k):27            if isinstance(k, tuple):28                r = SimpleChainedObjects(29                    # super() within a list comprehension needs explicit args.30                    [super(Axes.AxisDict, self).__getitem__(k1) for k1 in k])31                return r32            elif isinstance(k, slice):33                if k.start is None and k.stop is None and k.step is None:34                    return SimpleChainedObjects(list(self.values()))35                else:36                    raise ValueError("Unsupported slice")37            else:38                return dict.__getitem__(self, k)39 40        def __call__(self, *v, **kwargs):41            return maxes.Axes.axis(self.axes, *v, **kwargs)42 43    @property44    def axis(self):45        return self._axislines46 47    def clear(self):48        # docstring inherited49        super().clear()50        # Init axis artists.51        self._axislines = self.AxisDict(self)52        self._axislines.update(53            bottom=SimpleAxisArtist(self.xaxis, 1, self.spines["bottom"]),54            top=SimpleAxisArtist(self.xaxis, 2, self.spines["top"]),55            left=SimpleAxisArtist(self.yaxis, 1, self.spines["left"]),56            right=SimpleAxisArtist(self.yaxis, 2, self.spines["right"]))57 58 59class SimpleAxisArtist(Artist):60    def __init__(self, axis, axisnum, spine):61        self._axis = axis62        self._axisnum = axisnum63        self.line = spine64 65        if isinstance(axis, XAxis):66            self._axis_direction = ["bottom", "top"][axisnum-1]67        elif isinstance(axis, YAxis):68            self._axis_direction = ["left", "right"][axisnum-1]69        else:70            raise ValueError(71                f"axis must be instance of XAxis or YAxis, but got {axis}")72        super().__init__()73 74    @property75    def major_ticks(self):76        tickline = "tick%dline" % self._axisnum77        return SimpleChainedObjects([getattr(tick, tickline)78                                     for tick in self._axis.get_major_ticks()])79 80    @property81    def major_ticklabels(self):82        label = "label%d" % self._axisnum83        return SimpleChainedObjects([getattr(tick, label)84                                     for tick in self._axis.get_major_ticks()])85 86    @property87    def label(self):88        return self._axis.label89 90    def set_visible(self, b):91        self.toggle(all=b)92        self.line.set_visible(b)93        self._axis.set_visible(True)94        super().set_visible(b)95 96    def set_label(self, txt):97        self._axis.set_label_text(txt)98 99    def toggle(self, all=None, ticks=None, ticklabels=None, label=None):100 101        if all:102            _ticks, _ticklabels, _label = True, True, True103        elif all is not None:104            _ticks, _ticklabels, _label = False, False, False105        else:106            _ticks, _ticklabels, _label = None, None, None107 108        if ticks is not None:109            _ticks = ticks110        if ticklabels is not None:111            _ticklabels = ticklabels112        if label is not None:113            _label = label114 115        if _ticks is not None:116            tickparam = {f"tick{self._axisnum}On": _ticks}117            self._axis.set_tick_params(**tickparam)118        if _ticklabels is not None:119            tickparam = {f"label{self._axisnum}On": _ticklabels}120            self._axis.set_tick_params(**tickparam)121 122        if _label is not None:123            pos = self._axis.get_label_position()124            if (pos == self._axis_direction) and not _label:125                self._axis.label.set_visible(False)126            elif _label:127                self._axis.label.set_visible(True)128                self._axis.set_label_position(self._axis_direction)129