1from abc import ABC
  2from typing import Any, List, Optional, Union
  3from io import BytesIO
  4import base64
  5import matplotlib.pyplot as plt
  6import markdown
  7from IPython import display
  8
  9from fmeval.reporting.constants import CENTER, ListType, MARKDOWN_EXTENSIONS, RIGHT
 10from fmeval.reporting.constants import SINGLE_NEWLINE, DOUBLE_NEWLINE
 11
 12
 13class Cell(ABC):
 14    """
 15    Base class for a report cell.
 16    """
 17
 18
 19class MarkdownCell(Cell):
 20    """
 21    Base class representing a markdown cell.
 22    """
 23
 24    content: List[Union[str, "MarkdownCell"]]
 25
 26    def __init__(self, *args):
 27        """
 28        Input may be strings or MarkdownCells.
 29        Examples:
 30            assert str(MarkdownCell("# Hello1")) == "# Hello1"
 31            assert str(MarkdownCell("# Hello1", "# Hello2")) == "# Hello1  \n\n# Hello2"
 32            assert str(MarkdownCell(MarkdownCell("# Hello1"), "# Hello2") == "# Hello1  \n\n# Hello2"
 33        """
 34        assert all(isinstance(arg, str) or isinstance(arg, MarkdownCell) for arg in args)
 35        self.content = list(args)
 36
 37    def __str__(self):
 38        return DOUBLE_NEWLINE.join([str(s) for s in self.content])
 39
 40    def show(self):  # pragma: no cover
 41        """
 42        Displays the cell content in an IPython notebook cell.
 43        """
 44        content = markdown.markdown(self.__str__(), extensions=MARKDOWN_EXTENSIONS)
 45        return display.HTML(content)
 46
 47
 48class HeadingCell(MarkdownCell):
 49    """
 50    This class represents a Markdown heading.
 51    """
 52
 53    def __init__(self, text: str, level: int):
 54        """
 55        :param text: The text for this header
 56        :param level: The heading level
 57        """
 58        super().__init__(f"{'#' * level} {text}")
 59
 60
 61class BoldCell(MarkdownCell):
 62    """
 63    This class represents a bold piece of text.
 64    """
 65
 66    def __init__(self, text):
 67        super().__init__(f"**{text}**")
 68
 69
 70class ListCell(MarkdownCell):
 71    """
 72    Creates a bulleted or numbered list.
 73    """
 74
 75    def __init__(self, items: List[str], list_type: ListType):
 76        """
 77        :param items:  A list of strings where each string represents one item in the list.
 78        :param list_type: Whether the list is bulleted or numbered.
 79        """
 80        if list_type == ListType.NUMBERED:
 81            list_content = SINGLE_NEWLINE.join(f"{idx}. {value}" for idx, value in enumerate(items, start=1))
 82        else:
 83            list_content = SINGLE_NEWLINE.join(f"* {value}" for value in items)
 84        super().__init__(list_content)
 85
 86
 87class ColumnsLayoutCell(MarkdownCell):
 88    """
 89    This class creates a multi-column layout cell
 90    """
 91
 92    def __init__(self, columns: List[List[Any]]):
 93        """
 94        :param columns: A list of Lists of strings or MarkdownCells, where each inner list is one column
 95        """
 96        div_begin = '<div class="row" markdown="1">'
 97        div_end = "</div>"
 98        col_div = f'<div class="column" markdown="1" style="float: left;width: {str(int(100//len(columns)))}%;">\n'
 99        result = "".join(
100            [col_div + SINGLE_NEWLINE.join([str(item) for item in col]) + SINGLE_NEWLINE + div_end for col in columns]
101        )
102        super().__init__(div_begin, result, div_end, ' <br style="clear:both" />')
103
104
105class FigureCell(MarkdownCell):
106    """
107    This class represents a MarkdownCell containing HTML for a Pyplot Figure.
108    """
109
110    def __init__(
111        self,
112        fig: plt.Figure,
113        width: Optional[str] = None,
114        height: Optional[str] = None,
115        center: Optional[bool] = True,
116    ):
117        """
118        Initializes a FigureCell.
119
120        :param fig: The Pyplot figure that this cell represents.
121        :param width: See _html_wrapper docstring
122        :param height: See _html_wrapper docstring
123        :param center: if the figure is center aligned
124        """
125        encoded = FigureCell._encode(fig)
126        html = FigureCell._html_wrapper(encoded=encoded, height=height, width=width, center=center)
127        super().__init__(html)
128
129    @staticmethod
130    def _encode(fig: plt.Figure) -> bytes:
131        """
132        Returns the base64 encoding of `fig`.
133
134        :param fig: The Pyplot Figure to be encoded.
135        :returns: The base64 encoding of `fig`
136        """
137        buffer = BytesIO()
138        fig.tight_layout()
139        fig.savefig(buffer, format="png", bbox_inches="tight")
140        plt.close(fig)  # save memory by closing the figure
141        buffer.seek(0)
142        encoded = base64.b64encode(buffer.getvalue())
143        return encoded
144
145    @staticmethod
146    def _html_wrapper(
147        encoded: bytes, height: Optional[str], width: Optional[str], center: Optional[bool] = True
148    ) -> str:
149        """
150        Decodes the provided base64-encoded bytes (which will generally correspond to
151        an encoded Pyplot Figure) as a string, then wraps it in HTML.
152
153        :param encoded: The base64 encoded bytes to be wrapped
154        :param height: The `height` HTML attribute
155        :param width: The `width` HTML attribute
156        :param center: Horizontal centering of the figure
157        :returns: An HTML string representing the wrapped encoded string
158        """
159        style = 'style="display: block;'
160        if center:  # pragma: no branch
161            style += "margin-left:auto; margin-right: auto;"
162        if width:
163            style += f"width:{width}; "
164        if height:
165            style += f"height:{height};"
166        encoded_str = encoded.decode("utf-8")
167        html = f"<br><img {style}\" src='data:image/png;base64,{encoded_str}'><br>"
168        return html
169
170
171class BarPlotCell(FigureCell):
172    """
173    This class represents a Pyplot bar plot figure.
174    """
175
176    def __init__(
177        self,
178        labels: List[str],
179        heights: List[Union[int, float]],
180        color: Optional[Union[List[str], str]] = None,
181        title: str = "Title",
182        plot_height: Optional[str] = None,
183        plot_width: Optional[str] = None,
184        center: Optional[bool] = True,
185        origin: float = 0,
186    ):
187        """
188        Initializes a BarPlotCell.
189
190        :param labels: The labels corresponding to each of the bars in the plot
191        :param heights: The heights of the bars in the plot
192        :param title: The title of the bar plot
193        :param plot_height: Height of the plot as a string
194        :param plot_width: Width the plot as a string
195        :param center: Boolean indicating if the plot should be center aligned in the page
196        """
197        assert len(labels) == len(
198            heights
199        ), f"Number of labels in {labels} does not match number of bar heights in {heights}"
200        fig = BarPlotCell._create_bar_plot_fig(labels, heights, color=color, title=title, origin=origin)
201        super().__init__(fig, height=plot_height, width=plot_width, center=center)
202
203    @staticmethod
204    def _create_bar_plot_fig(
205        labels: List[str],
206        heights: List[Union[int, float]],
207        color: Optional[Union[List[str], str]] = None,
208        title: str = "Title",
209        set_spines_visible: bool = False,
210        set_ticks_visible: bool = False,
211        set_horizontal_grid_lines: bool = True,
212        max_bar_width: float = 0.3,
213        origin: float = 0,
214    ) -> plt.Figure:
215        fig, ax = plt.subplots()
216        heights = [height - origin for height in heights]
217        ax.bar(labels, heights, width=max_bar_width, color=color)
218        locs = ax.get_yticks()
219        ax.set_yticks(locs, [round(loc + origin, ndigits=3) for loc in locs])
220        if origin != 0:  # pragma: no cover
221            ax.axhline(0, color="gray")
222            ax.text(
223                x=1.02,
224                y=0,
225                s="unbiased model",
226                va="center",
227                ha="left",
228                bbox=dict(facecolor="w", alpha=0.5),
229                transform=ax.get_yaxis_transform(),
230            )
231        ax.set_title(title)
232
233        # auto-format bar labels to not overlap
234        fig.autofmt_xdate()
235        if set_horizontal_grid_lines:  # pragma: no branch
236            ax.grid(axis="y")
237            ax.set_axisbelow(True)
238        if not set_spines_visible:  # pragma: no branch
239            ax.spines["top"].set_visible(False)
240            ax.spines["right"].set_visible(False)
241            ax.spines["bottom"].set_visible(False)
242            ax.spines["left"].set_visible(False)
243        if not set_ticks_visible:  # pragma: no branch
244            plt.tick_params(axis="both", which="both", length=0)
245
246        return fig
247
248
249class TableCell(MarkdownCell):
250    """
251    This class represents an HTML table.
252
253    Note that despite having "Cell" in its name, this class does *not*
254    represent a single cell within an HTML table, but rather the entire table.
255    The "Cell" suffix is included to match the naming convention for subclasses
256    of MarkdownCell.
257    """
258
259    def __init__(
260        self,
261        data: List[List[Any]],
262        headers: List[str],
263        table_align: str = CENTER,
264        cell_align: str = RIGHT,
265        style: Optional[str] = None,
266        caption: Optional[str] = None,
267    ):
268        """
269        Initializes a TableCell.
270
271        :param data: A 2D array representing tabular data
272        :param headers: The table's headers, i.e. column names
273        :param table_align: The alignment of the table within the overarching markdown
274        :param cell_align: The alignment of text within each cell of the table
275        """
276        assert len(headers) == len(data[0]), (
277            f"Number of headers in {headers} does not match " f"the number of columns in the data: {len(data[0])}"
278        )
279        html = TableCell._create_table_html(data, headers, table_align, cell_align, style, caption)
280        super().__init__(html)
281
282    @staticmethod
283    def _create_table_html(
284        data: List[List[Any]],
285        headers: List[str],
286        table_align: str,
287        cell_align: str,
288        style: Optional[str],
289        caption: Optional[str],
290    ) -> str:
291        """
292        Creates the HTML for a table.
293
294        :param data: A 2D array representing tabular data
295        :param headers: The table's headers, i.e. column names
296        :param table_align: The alignment of the table within the overarching markdown
297        :param cell_align: The alignment of text within each cell of the table
298        :returns: A string encoding the HTML for the table
299        """
300        table_style = f'style="{style}"' if style else ""
301        html = [
302            f"<table align={table_align} {table_style}>",
303            f'<caption style="text-align: left; padding-bottom: 15px;">{caption}</caption>' if caption else "",
304            TableCell._create_table_row(headers, cell_align, is_header=True),
305        ]
306        for row in data:
307            html.append(TableCell._create_table_row(row, cell_align))
308        html.append("</table>")
309        return SINGLE_NEWLINE.join(html)
310
311    @staticmethod
312    def _create_table_row(row: List[Any], cell_align: str, is_header: bool = False) -> str:
313        """
314        Creates the HTML for a single table row.
315
316        :param row: A list representing the elements in the row
317        :param cell_align: The alignment of text within each cell of the table
318        :is_header: Whether `row` corresponds to the table header
319        :returns: A string encoding the HTML for the table row
320        """
321        tag = "th" if is_header else "td"
322        html = ["<tr>"]
323        for i, elem in enumerate(row):
324            style = f'style="text-align: {cell_align};"' if i != 0 else ""
325            html.append(f"<{tag} {style}>{elem}</{tag}>")
326        html.append("</tr>")
327        return " ".join(html)