diff --git a/edvart/report.py b/edvart/report.py index dd75018..5696dc0 100755 --- a/edvart/report.py +++ b/edvart/report.py @@ -45,14 +45,14 @@ def __init__( self.df = dataframe self.sections: list[Section] = [] self.verbosity = Verbosity(verbosity) + self._table_of_contents = None def show(self) -> None: """Renders the report in the calling notebook.""" + if self._table_of_contents is not None: + self._table_of_contents.show(self.sections) for section in self.sections: - if isinstance(section, TableOfContents): - section.show(self.sections) - else: - section.show(self.df) + section.show(self.df) def _select_columns( self, @@ -172,11 +172,10 @@ def _generate_notebook( nb["cells"].append(nbf4.new_code_cell(load_df)) # Generate code for each report section + if self._table_of_contents is not None: + self._table_of_contents.add_cells(self.sections, nb["cells"]) for section in self.sections: - if isinstance(section, TableOfContents): - section.add_cells(self.sections, nb["cells"]) - else: - section.add_cells(nb["cells"]) + section.add_cells(nb["cells"]) return nb @@ -591,7 +590,7 @@ def add_table_of_contents(self, include_subsections: bool = True) -> "ReportBase contents. However, they won't be included in an exported notebook created by report's export_notebook function. """ - self.sections.append(TableOfContents(include_subsections)) + self._table_of_contents = TableOfContents(include_subsections) return self diff --git a/edvart/report_sections/table_of_contents.py b/edvart/report_sections/table_of_contents.py index cf8aeb5..4aeea73 100644 --- a/edvart/report_sections/table_of_contents.py +++ b/edvart/report_sections/table_of_contents.py @@ -86,11 +86,8 @@ def add_cells(self, sections: List[Section], cells: List[Dict[str, Any]]) -> Non cells.append(section_header) lines: List[str] = [] - # Add links to all main sections (not including subsections) besides the first (table of - # content) section for section in sections: - if not isinstance(section, TableOfContents): - lines.append(TableOfContents._get_section_link(section, 1)) + lines.append(TableOfContents._get_section_link(section, 1)) cells.append(nbfv4.new_markdown_cell("\n".join(lines))) # pylint: disable=arguments-renamed @@ -105,9 +102,6 @@ def show(self, sections: List[Section]) -> None: display(Markdown(self.get_title(section_level=1))) lines = [] - # Add links to all sections including their subsections besides the first (table of content) - # section for section in sections: - if not isinstance(section, TableOfContents): - self._add_section_lines(section, 1, lines, self._include_subsections) + self._add_section_lines(section, 1, lines, self._include_subsections) display(Markdown("\n".join(lines))) diff --git a/tests/test_report.py b/tests/test_report.py index b40e48f..7d2cd67 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -40,14 +40,14 @@ def test_default_report(): ) assert len(report.sections) > 0, "Default report should not be empty" - assert report.sections[1].verbosity == Verbosity.MEDIUM, "Wrong section verbosity" - assert report.sections[1].columns is None, "Default column selection should be None" + assert report.sections[0].verbosity == Verbosity.MEDIUM, "Wrong section verbosity" + assert report.sections[0].columns is None, "Default column selection should be None" - assert report.sections[2].verbosity == Verbosity.HIGH, "Wrong section verbosity" - assert report.sections[2].columns is None, "Default column selection should be None" + assert report.sections[1].verbosity == Verbosity.HIGH, "Wrong section verbosity" + assert report.sections[1].columns is None, "Default column selection should be None" - assert report.sections[3].verbosity == Verbosity.LOW, "Wrong section verbosity" - assert report.sections[3].columns == ["Col1", "Col2", "Col3"], "Wrong columns" + assert report.sections[2].verbosity == Verbosity.LOW, "Wrong section verbosity" + assert report.sections[2].columns == ["Col1", "Col2", "Col3"], "Wrong columns" def test_column_selection():