fix: fix imports in exported code (#153)

datamole-ai · Sep 27, 2023 · 50d1c8a · 50d1c8a
1 parent 0ed449b
commit 50d1c8a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 11 deletions.
diff --git a/edvart/report_sections/multivariate_analysis.py b/edvart/report_sections/multivariate_analysis.py
@@ -527,7 +527,7 @@ def required_imports(self) -> List[str]:
             "from typing import Iterable",
             "import plotly",
             "import plotly.graph_objects as go",
-            "plotly.offline.init_notebook_mode()",
+            "from edvart.data_types import is_boolean, is_categorical, is_numeric",
         ]
 
     def add_cells(self, cells: List[Dict[str, Any]], df: pd.DataFrame) -> None:
@@ -705,9 +705,7 @@ def required_imports(self) -> List[str]:
             return ["from edvart.report_sections.multivariate_analysis import parallel_categories"]
         return [
             "from edvart.utils import discrete_colorscale",
-            "import plotly",
             "import plotly.graph_objects as go",
-            "plotly.offline.init_notebook_mode()",
         ]
 
     def add_cells(self, cells: List[Dict[str, Any]], df: pd.DataFrame) -> None:

diff --git a/edvart/report_sections/timeseries_analysis/autocorrelation.py b/edvart/report_sections/timeseries_analysis/autocorrelation.py
@@ -11,7 +11,7 @@
 
 from edvart.data_types import is_numeric
 from edvart.decorators import check_index_time_ascending
-from edvart.report_sections.code_string_formatting import get_code, total_dedent
+from edvart.report_sections.code_string_formatting import get_code
 from edvart.report_sections.section_base import Section, Verbosity
 
 
@@ -50,13 +50,8 @@ def required_imports(self) -> List[str]:
             ]
         if self.verbosity == Verbosity.MEDIUM:
             return [
-                total_dedent(
-                    """
-                    from edvart.report_sections.timeseries_analysis import Autocorrelation
-                    plot_acf = Autocorrelation.plot_acf
-                    plot_pacf = Autocorrelation.plot_pacf
-                    """
-                )
+                "from edvart.report_sections.timeseries_analysis.autocorrelation"
+                " import plot_acf, plot_pacf"
             ]
         # Verbosity.HIGH
         return [