diff --git a/src/vuecore/constants.py b/src/vuecore/constants.py index 0d0c4f6..2c2003e 100644 --- a/src/vuecore/constants.py +++ b/src/vuecore/constants.py @@ -3,15 +3,38 @@ try: from enum import StrEnum except ImportError: + # Fallback for Python versions < 3.11 that don't have StrEnum built-in from strenum import StrEnum class PlotType(StrEnum): + """Enum representing supported plot types.""" + SCATTER = auto() LINE = auto() BAR = auto() class EngineType(StrEnum): + """Enum representing supported plotting engines.""" + PLOTLY = auto() # Add other engines as needed + + +class OutputFileFormat(StrEnum): + """Enum representing supported output file formats.""" + + PNG = auto() + JPG = auto() + JPEG = auto() + SVG = auto() + PDF = auto() + HTML = auto() + JSON = auto() + WEBP = auto() + + @property + def value_with_dot(self): + """Return the file extension with the dot (e.g., '.png').""" + return f".{self.value}" diff --git a/src/vuecore/engines/plotly/saver.py b/src/vuecore/engines/plotly/saver.py index 654c202..f9887b0 100644 --- a/src/vuecore/engines/plotly/saver.py +++ b/src/vuecore/engines/plotly/saver.py @@ -2,6 +2,8 @@ import kaleido from pathlib import Path +from vuecore.constants import OutputFileFormat + def save(fig: go.Figure, filepath: str) -> None: """ @@ -33,34 +35,53 @@ def save(fig: go.Figure, filepath: str) -> None: suffix = path.suffix.lower() try: - if suffix in [".png", ".jpg", ".jpeg", ".webp", ".svg", ".pdf"]: + # Define static suffixes from the OutputFileFormat enum + image_suffixes = [ + OutputFileFormat.PNG.value_with_dot, + OutputFileFormat.JPG.value_with_dot, + OutputFileFormat.JPEG.value_with_dot, + OutputFileFormat.WEBP.value_with_dot, + OutputFileFormat.SVG.value_with_dot, + OutputFileFormat.PDF.value_with_dot, + ] + + if suffix in image_suffixes: try: fig.write_image(filepath) except RuntimeError as e: + # Handle specific Kaleido errors for Chrome installation if "Kaleido requires Google Chrome" in str(e): print( "[VueCore] Chrome not found. Attempting automatic install using `kaleido.get_chrome_sync()`..." ) try: kaleido.get_chrome_sync() - fig.write_image(filepath) # Retry after installing Chrome + # Retry after installing Chrome + fig.write_image(filepath) except Exception as install_error: raise RuntimeError( "[VueCore] Failed to install Chrome automatically. " "Please install it manually or run `plotly_get_chrome`." ) from install_error else: - raise - elif suffix == ".html": + raise # Re-raise other RuntimeError exceptions + elif suffix == OutputFileFormat.HTML.value_with_dot: fig.write_html(filepath, include_plotlyjs="cdn") - elif suffix == ".json": - fig.write_json(filepath, pretty=True) + elif suffix == OutputFileFormat.JSON.value_with_dot: + fig.write_json( + filepath, pretty=True + ) # Added pretty=True for readable JSON output else: + # Generate a dynamic list of supported formats for the error message + supported_suffixes = ", ".join( + [f"'{f.value_with_dot}'" for f in OutputFileFormat] + ) raise ValueError( f"Unsupported file format: '{suffix}'. " - "Supported formats: .png, .jpg, .jpeg, .webp, .svg, .pdf, .html, .json" + f"Supported formats: {supported_suffixes}" ) except Exception as e: + # Catch any exceptions during the saving process and re-raise as a RuntimeError raise RuntimeError(f"[VueCore] Failed to save plot: {filepath}") from e print(f"[VueCore] Plot saved to {filepath}")