feat: Export to JSON (#94)

qTipTip · web-flow · commit 29b5c0e1d095 · 2025-08-22T11:30:28.000+02:00
* Improve export to include JSON support

* Add JSON exports and associated tests

* Remove unused export mode enum

* Fix JSON format
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,18 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [Unreleased]
+
+### Added
+- JSON export functionality for palettes with `to_json()` method
+- General `export()` method supporting multiple formats (JSON, CSV) with JSON as default
+- CLI `--export-json` flag for JSON output instead of CSV to stdout
+- CLI `--output` parameter with auto-detection for individual vs combined file export
+- Smart file naming with prefix + index (`palette_001.json`) for directory exports
+- Support for combined JSON files containing multiple palettes
+- `hex` property on `Color` class, representing RGB-values in hexadecimal
+- Pylette JSON format with colorspace-specific field names (e.g., `rgb`, `hsv`, `hls`)
+
 # Released
 
 ## 4.4.0 - 20/08/2025
diff --git a/Pylette/cmd.py b/Pylette/cmd.py
@@ -1,3 +1,4 @@
+import json
 import pathlib
 from enum import Enum
 from typing import Annotated, List
@@ -6,20 +7,14 @@
 
 from Pylette.src.cli_utils import PyletteProgress
 from Pylette.src.color_extraction import batch_extract_colors
-from Pylette.src.types import BatchResult, ExtractionMethod
+from Pylette.src.types import BatchResult, ColorSpace, ExtractionMethod
 
 
 class SortBy(str, Enum):
     frequency = "frequency"
     luminance = "luminance"
 
 
-class ColorSpace(str, Enum):
-    rgb = "rgb"
-    hsv = "hsv"
-    hls = "hls"
-
-
 pylette_app = typer.Typer()
 
 
@@ -34,7 +29,7 @@ def main(
     stdout: bool = True,
     out_filename: pathlib.Path | None = None,
     display_colors: bool = False,
-    colorspace: ColorSpace = ColorSpace.rgb,
+    colorspace: ColorSpace = ColorSpace.RGB,
     alpha_mask_threshold: int | None = typer.Option(
         None,
         min=0,
@@ -44,7 +39,19 @@ def main(
     num_threads: int | None = typer.Option(
         None, min=1, help="Number of threads used for batch extraction of color palettes"
     ),
+    export_json: bool = typer.Option(False, "--export-json", help="Export palettes to JSON format"),
+    output: pathlib.Path | None = typer.Option(
+        None,
+        "--output",
+        help="Output file or directory for JSON export."
+        "If directory: creates individual files. If file: creates combined file.",
+    ),
 ):
+    # Validate export_json requirements
+    if export_json and output is None:
+        typer.echo("Error: --output is required when using --export-json", err=True)
+        raise typer.Exit(1)
+
     output_file_path = str(out_filename) if out_filename is not None else None
 
     # Set up progress bar for CLI
@@ -74,12 +81,19 @@ def progress_callback(task_number: int, result: BatchResult):
 
     successful = [r for r in results if r.success]
     failed = [r for r in results if not r.success]
-    for success in successful:
-        if success.palette is not None:
-            success.palette.to_csv(
-                filename=output_file_path, frequency=True, stdout=stdout, colorspace=colorspace.value
-            )
-            if display_colors:
+
+    if export_json and output:
+        handle_json_export(successful, output, colorspace)
+    else:
+        # Original CSV behavior
+        for success in successful:
+            if success.palette is not None:
+                success.palette.to_csv(filename=output_file_path, frequency=True, stdout=stdout, colorspace=colorspace)
+
+    # Display colors if requested
+    if display_colors:
+        for success in successful:
+            if success.palette is not None:
                 success.palette.display()
 
     if failed:
@@ -93,6 +107,50 @@ def progress_callback(task_number: int, result: BatchResult):
         raise typer.Exit(2)
 
 
+def handle_json_export(
+    successful_results: list[BatchResult], output_path: pathlib.Path, colorspace: ColorSpace
+) -> None:
+    """Handle JSON export for successful palette extractions."""
+
+    if output_path.is_dir() or (not output_path.exists() and not output_path.suffix):
+        # Directory mode: individual files
+        output_dir = output_path
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        for i, result in enumerate(successful_results, 1):
+            if result.palette is not None:
+                filename = output_dir / f"palette_{i:03d}.json"
+                result.palette.to_json(filename=str(filename), colorspace=colorspace, stdout=False)
+
+        typer.echo(f"✓ Exported {len(successful_results)} palettes to {output_dir}")
+
+    else:
+        # File mode: combined file
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Create combined JSON structure
+        combined_data: dict[str, object] = {
+            "palettes": [],
+            "total_count": len(successful_results),
+            "colorspace": colorspace,
+        }
+        palettes_list = []
+
+        for result in successful_results:
+            if result.palette is not None:
+                palette_data = result.palette.to_json(filename=None, colorspace=colorspace, stdout=False)
+                if palette_data is not None:
+                    palettes_list.append(palette_data)
+
+        combined_data["palettes"] = palettes_list
+
+        # Write combined file
+        with open(output_path, "w") as f:
+            json.dump(combined_data, f, indent=2)
+
+        typer.echo(f"✓ Exported {len(successful_results)} palettes to {output_path}")
+
+
 def print_extraction_summary(successful: list[BatchResult], failed: list[BatchResult]):
     total = len(successful) + len(failed)
 
diff --git a/Pylette/src/color.py b/Pylette/src/color.py
@@ -1,8 +1,10 @@
 import colorsys
-from typing import Literal, cast
+from typing import cast
 
 import numpy as np
 
+from Pylette.src.types import ColorSpace
+
 # Weights for calculating luminance
 luminance_weights = np.array([0.2126, 0.7152, 0.0722])
 
@@ -50,17 +52,17 @@ def __lt__(self, other: "Color") -> bool:
         """
         return self.freq < other.freq
 
-    def get_colors(self, colorspace: Literal["rgb", "hsv", "hls"] = "rgb") -> tuple[int, ...] | tuple[float, ...]:
+    def get_colors(self, colorspace: ColorSpace = ColorSpace.RGB) -> tuple[int, ...] | tuple[float, ...]:
         """
         Returns the color values in the specified color space.
 
         Parameters:
-            colorspace (Literal["rgb", "hsv", "hls"]): The color space to use.
+            colorspace (ColorSpace): The color space to use.
 
         Returns:
             tuple[int, ...] | tuple[float, ...]: The color values in the specified color space.
         """
-        colors = {"rgb": self.rgb, "hsv": self.hsv, "hls": self.hls}
+        colors = {ColorSpace.RGB: self.rgb, ColorSpace.HSV: self.hsv, ColorSpace.HLS: self.hls}
         return colors[colorspace]
 
     @property
@@ -83,6 +85,16 @@ def hls(self) -> tuple[float, float, float]:
         """
         return colorsys.rgb_to_hls(r=self.rgb[0] / 255, g=self.rgb[1] / 255, b=self.rgb[2] / 255)
 
+    @property
+    def hex(self) -> str:
+        """
+        Returns the color as a hexadecimal string.
+
+        Returns:
+            str: The color in hexadecimal format (e.g., "#FF5733").
+        """
+        return f"#{self.rgb[0]:02X}{self.rgb[1]:02X}{self.rgb[2]:02X}"
+
     @property
     def luminance(self) -> float:
         """
diff --git a/Pylette/src/palette.py b/Pylette/src/palette.py
@@ -1,10 +1,11 @@
+import json
 from typing import Literal
 
 import numpy as np
 from PIL import Image
 
 from Pylette.src.color import Color
-from Pylette.src.types import ExtractionParams, ImageInfo, PaletteMetaData, ProcessingStats, SourceType
+from Pylette.src.types import ColorSpace, ExtractionParams, ImageInfo, PaletteMetaData, ProcessingStats, SourceType
 
 
 class Palette:
@@ -92,7 +93,7 @@ def to_csv(
         self,
         filename: str | None = None,
         frequency: bool = True,
-        colorspace: Literal["rgb", "hsv", "hls"] = "rgb",
+        colorspace: ColorSpace = ColorSpace.RGB,
         stdout: bool = True,
     ):
         """
@@ -143,6 +144,125 @@ def random_color(self, N: int, mode: str = "frequency") -> list[Color]:
         else:
             raise ValueError(f"Invalid mode: {mode}. Must be 'frequency' or 'uniform'.")
 
+    def to_json(
+        self,
+        filename: str | None = None,
+        colorspace: ColorSpace = ColorSpace.RGB,
+        include_metadata: bool = True,
+        stdout: bool = True,
+    ) -> dict[str, object] | None:
+        """
+        Exports the palette to JSON format.
+
+        Parameters:
+            filename (str | None): File to save to. If None, returns the dictionary.
+            colorspace (Literal["rgb", "hsv", "hls"]): Color space to use.
+            include_metadata (bool): Whether to include palette metadata.
+            stdout (bool): Whether to print to stdout.
+
+        Returns:
+            dict | None: The palette data as a dictionary if filename is None.
+        """
+
+        # Build the palette data
+        palette_data: dict[str, object] = {
+            "colors": [],
+            "palette_size": self.number_of_colors,
+            "colorspace": colorspace,
+        }
+
+        colors_list = []
+        # Add color data
+        for color in self.colors:
+            color_values = color.get_colors(colorspace)
+            color_data: dict[str, object] = {
+                "frequency": float(color.freq),
+            }
+
+            # Add colorspace-specific field
+            colorspace_field = colorspace.value.lower()  # "rgb", "hsv", "hls"
+            if colorspace == ColorSpace.RGB:
+                # RGB values should be integers
+                color_data[colorspace_field] = [int(v) if isinstance(v, np.integer) else v for v in color_values]
+            else:
+                # HSV/HLS values should be floats
+                color_data[colorspace_field] = [
+                    float(v) if isinstance(v, (np.integer, np.floating)) else v for v in color_values
+                ]
+
+            # Add hex (always present, derived from RGB)
+            color_data["hex"] = color.hex
+
+            # Add RGB reference if colorspace is not RGB
+            if colorspace != ColorSpace.RGB:
+                color_data["rgb"] = [int(v) if isinstance(v, np.integer) else v for v in color.rgb]
+
+            colors_list.append(color_data)
+
+        palette_data["colors"] = colors_list
+
+        # Add metadata if requested and available
+        if include_metadata and self.metadata:
+            metadata_dict: dict[str, object] = {}
+
+            if "image_source" in self.metadata:
+                metadata_dict["image_source"] = self.metadata["image_source"]
+            if "source_type" in self.metadata:
+                metadata_dict["source_type"] = self.metadata["source_type"]
+            if "extraction_params" in self.metadata:
+                metadata_dict["extraction_params"] = self.metadata["extraction_params"]
+            if "image_info" in self.metadata:
+                metadata_dict["image_info"] = self.metadata["image_info"]
+            if "processing_stats" in self.metadata:
+                metadata_dict["processing_stats"] = self.metadata["processing_stats"]
+
+            palette_data["metadata"] = metadata_dict
+
+        # Print to stdout if requested
+        if stdout:
+            print(json.dumps(palette_data, indent=2))
+
+        # Save to file if filename provided
+        if filename is not None:
+            with open(filename, "w") as f:
+                json.dump(palette_data, f, indent=2)
+            return None
+
+        # Return data if no filename provided
+        return palette_data
+
+    def export(
+        self,
+        filename: str,
+        format: Literal["json", "csv"] = "json",
+        colorspace: ColorSpace = ColorSpace.RGB,
+        include_frequency: bool = True,
+        include_metadata: bool = True,
+        stdout: bool = False,
+    ) -> None:
+        """
+        General export method that supports multiple formats with JSON as default.
+
+        Parameters:
+            filename (str): File to save to (extension will be added automatically).
+            format (Literal["json", "csv"]): Export format (default: json).
+            colorspace (Literal["rgb", "hsv", "hls"]): Color space to use.
+            include_frequency (bool): Whether to include frequency data.
+            include_metadata (bool): Whether to include metadata (JSON only).
+            stdout (bool): Whether to print to stdout.
+        """
+
+        if format == "json":
+            json_filename = f"{filename}.json"
+            self.to_json(
+                filename=json_filename, colorspace=colorspace, include_metadata=include_metadata, stdout=stdout
+            )
+        elif format == "csv":
+            csv_filename = f"{filename}.csv"
+            self.to_csv(filename=csv_filename, frequency=include_frequency, colorspace=colorspace, stdout=stdout)
+        else:
+            raise ValueError(f"Unsupported format: {format}. Supported formats: 'json', 'csv'")
+
     def __str__(self):
         return "".join(["({}, {}, {}, {}) \n".format(c.rgb[0], c.rgb[1], c.rgb[2], c.freq) for c in self.colors])
 
diff --git a/Pylette/src/types.py b/Pylette/src/types.py
@@ -58,6 +58,12 @@ class ExtractionMethod(str, Enum):
     KM = "KMeans"
 
 
+class ColorSpace(str, Enum):
+    RGB = "rgb"
+    HSV = "hsv"
+    HLS = "hls"
+
+
 # PaletteMetaData Types
 class SourceType(str, Enum):
     FILE_PATH = "file_path"
diff --git a/tests/integration/test_json_export.py b/tests/integration/test_json_export.py
diff --git a/tests/integration/test_metadata.py b/tests/integration/test_metadata.py
