pltr-cli 0.1.0__py3-none-any.whl

pltr/utils/formatting.py

@@ -0,0 +1,331 @@
+"""
+Output formatting utilities for CLI commands.
+"""
+
+import json
+import csv
+from typing import Any, Dict, List, Optional, Union
+from datetime import datetime
+from io import StringIO
+
+from rich.console import Console
+from rich.table import Table
+from rich import print as rich_print
+
+
+class OutputFormatter:
+    """Handles different output formats for CLI commands."""
+
+    def __init__(self, console: Optional[Console] = None):
+        """
+        Initialize formatter.
+
+        Args:
+            console: Rich console instance (creates one if not provided)
+        """
+        self.console = console or Console()
+
+    def format_output(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]]],
+        format_type: str = "table",
+        output_file: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Format data according to specified format.
+
+        Args:
+            data: Data to format
+            format_type: Output format ('table', 'json', 'csv')
+            output_file: Optional file path to write output
+
+        Returns:
+            Formatted string if no output file specified
+        """
+        if format_type == "json":
+            return self._format_json(data, output_file)
+        elif format_type == "csv":
+            return self._format_csv(data, output_file)
+        elif format_type == "table":
+            return self._format_table(data, output_file)
+        else:
+            raise ValueError(f"Unsupported format type: {format_type}")
+
+    def _format_json(
+        self, data: Any, output_file: Optional[str] = None
+    ) -> Optional[str]:
+        """Format data as JSON."""
+        # Convert datetime objects to strings for JSON serialization
+        data_serializable = self._make_json_serializable(data)
+        json_str = json.dumps(data_serializable, indent=2, default=str)
+
+        if output_file:
+            with open(output_file, "w") as f:
+                f.write(json_str)
+            return None
+        else:
+            rich_print(json_str)
+            return json_str
+
+    def _format_csv(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]]],
+        output_file: Optional[str] = None,
+    ) -> Optional[str]:
+        """Format data as CSV."""
+        if isinstance(data, dict):
+            data = [data]
+
+        if not data:
+            csv_str = ""
+        else:
+            # Get all unique keys for the CSV header
+            fieldnames_set: set[str] = set()
+            for item in data:
+                fieldnames_set.update(item.keys())
+            fieldnames = sorted(fieldnames_set)
+
+            output = StringIO()
+            writer = csv.DictWriter(output, fieldnames=fieldnames)
+            writer.writeheader()
+
+            for item in data:
+                # Convert complex objects to strings
+                row = {}
+                for key in fieldnames:
+                    value = item.get(key)
+                    if isinstance(value, (dict, list)):
+                        row[key] = json.dumps(value)
+                    elif value is None:
+                        row[key] = ""
+                    else:
+                        row[key] = str(value)
+                writer.writerow(row)
+
+            csv_str = output.getvalue()
+
+        if output_file:
+            with open(output_file, "w") as f:
+                f.write(csv_str)
+            return None
+        else:
+            print(csv_str, end="")
+            return csv_str
+
+    def _format_table(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]]],
+        output_file: Optional[str] = None,
+    ) -> Optional[str]:
+        """Format data as a rich table."""
+        if isinstance(data, dict):
+            data = [data]
+
+        if not data:
+            if output_file:
+                with open(output_file, "w") as f:
+                    f.write("No data to display\n")
+                return None
+            else:
+                self.console.print("No data to display")
+                return "No data to display"
+
+        # Create table
+        table = Table(show_header=True, header_style="bold blue")
+
+        # Get all unique columns
+        columns_set: set[str] = set()
+        for item in data:
+            columns_set.update(item.keys())
+        columns = sorted(columns_set)
+
+        # Add columns to table
+        for column in columns:
+            table.add_column(column, overflow="fold")
+
+        # Add rows
+        for item in data:
+            row = []
+            for column in columns:
+                value = item.get(column)
+                if isinstance(value, (dict, list)):
+                    # Format complex objects as JSON
+                    row.append(json.dumps(value, indent=2))
+                elif value is None:
+                    row.append("")
+                elif isinstance(value, datetime):
+                    row.append(value.isoformat())
+                else:
+                    row.append(str(value))
+            table.add_row(*row)
+
+        if output_file:
+            # For file output, convert to plain text
+            with open(output_file, "w") as f:
+                console = Console(file=f, force_terminal=False)
+                console.print(table)
+            return None
+        else:
+            self.console.print(table)
+            return str(table)
+
+    def _make_json_serializable(self, data: Any) -> Any:
+        """Convert data to JSON-serializable format."""
+        if isinstance(data, dict):
+            return {k: self._make_json_serializable(v) for k, v in data.items()}
+        elif isinstance(data, list):
+            return [self._make_json_serializable(item) for item in data]
+        elif isinstance(data, datetime):
+            return data.isoformat()
+        else:
+            return data
+
+    def format_dataset_list(
+        self,
+        datasets: List[Dict[str, Any]],
+        format_type: str = "table",
+        output_file: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Format dataset list with specific columns.
+
+        Args:
+            datasets: List of dataset dictionaries
+            format_type: Output format
+            output_file: Optional output file path
+
+        Returns:
+            Formatted string if no output file specified
+        """
+        # Select and order key columns for dataset display
+        formatted_datasets = []
+        for dataset in datasets:
+            formatted_dataset = {
+                "RID": dataset.get("rid", ""),
+                "Name": dataset.get("name", ""),
+                "Created": self._format_datetime(dataset.get("created_time")),
+                "Size": self._format_file_size(dataset.get("size_bytes")),
+                "Description": dataset.get("description", "")[:50] + "..."
+                if dataset.get("description", "")
+                else "",
+            }
+            formatted_datasets.append(formatted_dataset)
+
+        return self.format_output(formatted_datasets, format_type, output_file)
+
+    def format_dataset_detail(
+        self,
+        dataset: Dict[str, Any],
+        format_type: str = "table",
+        output_file: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Format detailed dataset information.
+
+        Args:
+            dataset: Dataset dictionary
+            format_type: Output format
+            output_file: Optional output file path
+
+        Returns:
+            Formatted string if no output file specified
+        """
+        if format_type == "table":
+            # For table format, show key-value pairs (only show fields that exist)
+            details = []
+
+            if dataset.get("rid"):
+                details.append({"Property": "RID", "Value": dataset["rid"]})
+            if dataset.get("name"):
+                details.append({"Property": "Name", "Value": dataset["name"]})
+            if dataset.get("parent_folder_rid"):
+                details.append(
+                    {"Property": "Parent Folder", "Value": dataset["parent_folder_rid"]}
+                )
+
+            # Add any other fields that might exist
+            for key, value in dataset.items():
+                if (
+                    key not in ["rid", "name", "parent_folder_rid"]
+                    and value is not None
+                    and value != ""
+                ):
+                    details.append(
+                        {"Property": key.replace("_", " ").title(), "Value": str(value)}
+                    )
+
+            return self.format_output(details, format_type, output_file)
+        else:
+            return self.format_output(dataset, format_type, output_file)
+
+    def format_file_list(
+        self,
+        files: List[Dict[str, Any]],
+        format_type: str = "table",
+        output_file: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Format file list with specific columns.
+
+        Args:
+            files: List of file dictionaries
+            format_type: Output format
+            output_file: Optional output file path
+
+        Returns:
+            Formatted string if no output file specified
+        """
+        # Format files for display
+        formatted_files = []
+        for file in files:
+            formatted_file = {
+                "Path": file.get("path", ""),
+                "Size": self._format_file_size(file.get("size_bytes")),
+                "Last Modified": self._format_datetime(file.get("last_modified")),
+                "Transaction": file.get("transaction_rid", "")[:12] + "..."
+                if file.get("transaction_rid")
+                else "",
+            }
+            formatted_files.append(formatted_file)
+
+        return self.format_output(formatted_files, format_type, output_file)
+
+    def _format_datetime(self, dt: Any) -> str:
+        """Format datetime for display."""
+        if dt is None:
+            return ""
+        if isinstance(dt, str):
+            return dt
+        if isinstance(dt, datetime):
+            return dt.strftime("%Y-%m-%d %H:%M:%S")
+        return str(dt)
+
+    def _format_file_size(self, size_bytes: Optional[int]) -> str:
+        """Format file size in human-readable format."""
+        if size_bytes is None:
+            return ""
+
+        if size_bytes < 1024:
+            return f"{size_bytes} B"
+        elif size_bytes < 1024**2:
+            return f"{size_bytes / 1024:.1f} KB"
+        elif size_bytes < 1024**3:
+            return f"{size_bytes / (1024**2):.1f} MB"
+        else:
+            return f"{size_bytes / (1024**3):.1f} GB"
+
+    def print_success(self, message: str):
+        """Print success message with formatting."""
+        self.console.print(f"✅ {message}", style="green")
+
+    def print_error(self, message: str):
+        """Print error message with formatting."""
+        self.console.print(f"❌ {message}", style="red")
+
+    def print_warning(self, message: str):
+        """Print warning message with formatting."""
+        self.console.print(f"⚠️  {message}", style="yellow")
+
+    def print_info(self, message: str):
+        """Print info message with formatting."""
+        self.console.print(f"ℹ️  {message}", style="blue")

pltr/utils/progress.py

@@ -0,0 +1,328 @@
+"""
+Progress bar utilities for long-running operations.
+"""
+
+from typing import Optional, Iterator, Any, Union
+from pathlib import Path
+from contextlib import contextmanager
+
+from rich.progress import (
+    Progress,
+    TextColumn,
+    BarColumn,
+    TaskProgressColumn,
+    TimeRemainingColumn,
+    FileSizeColumn,
+    TotalFileSizeColumn,
+    TransferSpeedColumn,
+    SpinnerColumn,
+)
+
+
+class FileProgressTracker:
+    """Progress tracker for file operations."""
+
+    def __init__(self, show_speed: bool = True):
+        """
+        Initialize progress tracker.
+
+        Args:
+            show_speed: Whether to show transfer speed
+        """
+        self.show_speed = show_speed
+        self._progress: Optional[Progress] = None
+
+    @contextmanager
+    def track_upload(
+        self, file_path: Union[str, Path], description: Optional[str] = None
+    ) -> Iterator[Any]:
+        """
+        Context manager for tracking file upload progress.
+
+        Args:
+            file_path: Path to file being uploaded
+            description: Optional description for progress bar
+
+        Yields:
+            Progress update function
+        """
+        file_path = Path(file_path)
+        total_size = file_path.stat().st_size
+        description = description or f"Uploading {file_path.name}"
+
+        columns = [
+            TextColumn("[bold blue]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            FileSizeColumn(),
+            TotalFileSizeColumn(),
+            TimeRemainingColumn(),
+        ]
+
+        if self.show_speed:
+            columns.append(TransferSpeedColumn())
+
+        with Progress(*columns) as progress:
+            self._progress = progress
+            task_id = progress.add_task(description, total=total_size)
+
+            def update_progress(bytes_transferred: int):
+                """Update progress with bytes transferred."""
+                progress.update(task_id, completed=bytes_transferred)
+
+            try:
+                yield update_progress
+            finally:
+                self._progress = None
+
+    @contextmanager
+    def track_download(
+        self,
+        target_path: Union[str, Path],
+        total_size: Optional[int] = None,
+        description: Optional[str] = None,
+    ) -> Iterator[Any]:
+        """
+        Context manager for tracking file download progress.
+
+        Args:
+            target_path: Path where file will be saved
+            total_size: Total file size in bytes (if known)
+            description: Optional description for progress bar
+
+        Yields:
+            Progress update function
+        """
+        target_path = Path(target_path)
+        description = description or f"Downloading {target_path.name}"
+
+        columns = [
+            TextColumn("[bold green]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+        ]
+
+        if total_size:
+            columns.extend(
+                [
+                    FileSizeColumn(),
+                    TotalFileSizeColumn(),
+                    TimeRemainingColumn(),
+                ]
+            )
+            if self.show_speed:
+                columns.append(TransferSpeedColumn())
+
+        with Progress(*columns) as progress:
+            self._progress = progress
+            task_id = progress.add_task(description, total=total_size)
+
+            def update_progress(bytes_transferred: int):
+                """Update progress with bytes transferred."""
+                progress.update(task_id, completed=bytes_transferred)
+
+            try:
+                yield update_progress
+            finally:
+                self._progress = None
+
+    @contextmanager
+    def track_operation(
+        self, description: str, total: Optional[int] = None
+    ) -> Iterator[Any]:
+        """
+        Context manager for tracking general operations.
+
+        Args:
+            description: Description of the operation
+            total: Total number of items (if known)
+
+        Yields:
+            Progress update function
+        """
+        columns = [
+            TextColumn("[bold yellow]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+        ]
+
+        if total:
+            columns.append(TimeRemainingColumn())
+
+        with Progress(*columns) as progress:
+            self._progress = progress
+            task_id = progress.add_task(description, total=total)
+
+            def update_progress(completed: int):
+                """Update progress with completed items."""
+                progress.update(task_id, completed=completed)
+
+            try:
+                yield update_progress
+            finally:
+                self._progress = None
+
+
+class SpinnerProgressTracker:
+    """Simple spinner for indeterminate operations."""
+
+    def __init__(self):
+        """Initialize spinner tracker."""
+        self._progress: Optional[Progress] = None
+
+    @contextmanager
+    def track_spinner(self, description: str) -> Iterator[None]:
+        """
+        Context manager for showing a spinner during operations.
+
+        Args:
+            description: Description of the operation
+
+        Yields:
+            None (operation runs in context)
+        """
+        columns = [
+            SpinnerColumn(),
+            TextColumn("[bold cyan]{task.description}"),
+        ]
+
+        with Progress(*columns, transient=True) as progress:
+            self._progress = progress
+            progress.add_task(description)
+
+            try:
+                yield
+            finally:
+                self._progress = None
+
+
+def create_file_chunks(
+    file_path: Union[str, Path], chunk_size: int = 8192
+) -> Iterator[bytes]:
+    """
+    Create file chunks for streaming upload with progress tracking.
+
+    Args:
+        file_path: Path to file to read
+        chunk_size: Size of each chunk in bytes
+
+    Yields:
+        File chunks as bytes
+    """
+    file_path = Path(file_path)
+    with open(file_path, "rb") as f:
+        while True:
+            chunk = f.read(chunk_size)
+            if not chunk:
+                break
+            yield chunk
+
+
+class ProgressCallbackAdapter:
+    """Adapter to convert progress callbacks to different formats."""
+
+    def __init__(self, update_callback, total_size: Optional[int] = None):
+        """
+        Initialize adapter.
+
+        Args:
+            update_callback: Callback function to update progress
+            total_size: Total size for percentage calculations
+        """
+        self.update_callback = update_callback
+        self.total_size = total_size
+        self.bytes_transferred = 0
+
+    def __call__(self, chunk_size: int):
+        """
+        Update progress with new chunk.
+
+        Args:
+            chunk_size: Size of transferred chunk
+        """
+        self.bytes_transferred += chunk_size
+        self.update_callback(self.bytes_transferred)
+
+    def get_percentage(self) -> Optional[float]:
+        """
+        Get current progress percentage.
+
+        Returns:
+            Percentage complete (0-100) or None if total size unknown
+        """
+        if self.total_size and self.total_size > 0:
+            return (self.bytes_transferred / self.total_size) * 100
+        return None
+
+
+# Utility functions for common progress patterns
+
+
+def with_upload_progress(
+    file_path: Union[str, Path], description: Optional[str] = None
+) -> FileProgressTracker:
+    """
+    Create a progress tracker configured for file uploads.
+
+    Args:
+        file_path: Path to file being uploaded
+        description: Optional description override
+
+    Returns:
+        Configured FileProgressTracker
+    """
+    return FileProgressTracker(show_speed=True)
+
+
+def with_download_progress(show_speed: bool = True) -> FileProgressTracker:
+    """
+    Create a progress tracker configured for file downloads.
+
+    Args:
+        show_speed: Whether to show transfer speed
+
+    Returns:
+        Configured FileProgressTracker
+    """
+    return FileProgressTracker(show_speed=show_speed)
+
+
+def with_spinner(description: str = "Working...") -> SpinnerProgressTracker:
+    """
+    Create a spinner for indeterminate operations.
+
+    Args:
+        description: Description of the operation
+
+    Returns:
+        Configured SpinnerProgressTracker
+    """
+    return SpinnerProgressTracker()
+
+
+# Example usage patterns
+"""
+# File upload with progress
+tracker = with_upload_progress("data.csv")
+with tracker.track_upload("data.csv") as progress:
+    # Upload file chunks
+    for chunk in create_file_chunks("data.csv"):
+        # ... upload chunk ...
+        progress(len(chunk))
+
+# File download with progress  
+tracker = with_download_progress()
+with tracker.track_download("output.csv", total_size=1024000) as progress:
+    # Download file
+    bytes_downloaded = 0
+    while not complete:
+        # ... download chunk ...
+        bytes_downloaded += chunk_size
+        progress(bytes_downloaded)
+
+# Indeterminate operation
+spinner = with_spinner("Processing dataset...")
+with spinner.track_spinner("Processing dataset..."):
+    # ... long running operation ...
+    pass
+"""