ai-coding-assistant 0.5.0__py3-none-any.whl

coding_assistant/documentation/generators/dataflow_generator.py

@@ -0,0 +1,440 @@
+"""Generate data flow visualizations.
+
+This module provides data flow diagram generation showing how data
+entities move through the system and are transformed by functions.
+"""
+
+from typing import List, Dict, Set, Optional
+from pathlib import Path
+import re
+
+from coding_assistant.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class DataFlowGenerator:
+    """
+    Generate data flow diagrams in Mermaid format.
+
+    Shows:
+    - Data entities (models, DTOs, data structures)
+    - Transformations (functions that process data)
+    - Data movement through the system
+    """
+
+    def __init__(self):
+        """Initialize the data flow generator."""
+        pass
+
+    def generate_dataflow_diagram(self,
+                                    data_entities: List[Dict],
+                                    transformations: List[Dict],
+                                    external_systems: Optional[List[Dict]] = None) -> str:
+        """
+        Generate data flow diagram in Mermaid format.
+
+        Args:
+            data_entities: List of data entities with structure:
+                          [{'name': 'User', 'type': 'model', 'file': 'models.py'}]
+            transformations: List of functions that transform data:
+                           [{'name': 'process_user', 'inputs': ['User'],
+                             'outputs': ['ProcessedUser']}]
+            external_systems: Optional list of external systems:
+                            [{'name': 'Database', 'type': 'storage'}]
+
+        Returns:
+            Mermaid flowchart as string
+        """
+        logger.info(f"Generating dataflow diagram with {len(data_entities)} entities, "
+                   f"{len(transformations)} transformations")
+
+        mermaid = ["flowchart LR"]
+
+        # Track all entity names
+        entity_names = {entity['name'] for entity in data_entities}
+
+        # Add external systems first (if any)
+        if external_systems:
+            for system in external_systems:
+                system_id = self._sanitize_id(system['name'])
+                system_type = system.get('type', 'external')
+
+                if system_type == 'storage':
+                    # Database/storage shape
+                    mermaid.append(f"    {system_id}[({system['name']})]")
+                else:
+                    # General external system
+                    mermaid.append(f"    {system_id}[/{system['name']}/]")
+
+        # Add data entities
+        for entity in data_entities:
+            entity_id = self._sanitize_id(entity['name'])
+            entity_type = entity.get('type', 'data')
+
+            # Different shapes based on entity type
+            if entity_type == 'model' or entity_type == 'class':
+                # Rounded rectangle for models
+                mermaid.append(f"    {entity_id}(({entity['name']}))")
+            elif entity_type == 'dto' or entity_type == 'interface':
+                # Hexagon for DTOs/interfaces
+                mermaid.append(f"    {entity_id}{{{{{entity['name']}}}}}")
+            else:
+                # Stadium shape for general data
+                mermaid.append(f"    {entity_id}([{entity['name']}])")
+
+        # Add transformations
+        for transform in transformations:
+            transform_id = self._sanitize_id(transform['name'])
+
+            # Rectangle for transformations (functions/methods)
+            mermaid.append(f"    {transform_id}[{transform['name']}]")
+
+            # Connect inputs to transformation
+            inputs = transform.get('inputs', [])
+            for input_entity in inputs:
+                input_id = self._sanitize_id(input_entity)
+
+                # Only connect if entity exists
+                if input_entity in entity_names or self._is_external_system(input_entity, external_systems):
+                    mermaid.append(f"    {input_id} --> {transform_id}")
+
+            # Connect transformation to outputs
+            outputs = transform.get('outputs', [])
+            for output_entity in outputs:
+                output_id = self._sanitize_id(output_entity)
+
+                # Only connect if entity exists
+                if output_entity in entity_names or self._is_external_system(output_entity, external_systems):
+                    # Dotted line for outputs to show creation/transformation
+                    mermaid.append(f"    {transform_id} -.-> {output_id}")
+
+        # Add styling
+        mermaid.extend(self._add_dataflow_styling(data_entities, transformations, external_systems))
+
+        result = "\n".join(mermaid)
+        logger.debug(f"Dataflow diagram generated ({len(mermaid)} lines)")
+
+        return result
+
+    def generate_pipeline_diagram(self,
+                                    pipeline_name: str,
+                                    stages: List[Dict]) -> str:
+        """
+        Generate data pipeline diagram showing sequential processing stages.
+
+        Args:
+            pipeline_name: Name of the pipeline
+            stages: List of pipeline stages with structure:
+                   [{'name': 'Extract', 'input': 'RawData', 'output': 'CleanData',
+                     'operations': ['validate', 'clean']}]
+
+        Returns:
+            Mermaid flowchart as string
+        """
+        logger.info(f"Generating pipeline diagram: {pipeline_name} with {len(stages)} stages")
+
+        mermaid = ["flowchart LR"]
+
+        # Add title as a subgraph
+        mermaid.append(f"    subgraph {self._sanitize_id(pipeline_name)}[{pipeline_name}]")
+
+        prev_output = None
+
+        for i, stage in enumerate(stages):
+            stage_id = self._sanitize_id(f"{pipeline_name}_{stage['name']}")
+
+            # Add stage as process node
+            mermaid.append(f"        {stage_id}[{stage['name']}]")
+
+            # Connect to previous stage
+            if i > 0 and prev_output:
+                data_id = self._sanitize_id(f"data_{i}")
+                mermaid.append(f"        {data_id}([{prev_output}])")
+                prev_stage_id = self._sanitize_id(f"{pipeline_name}_{stages[i-1]['name']}")
+                mermaid.append(f"        {prev_stage_id} --> {data_id}")
+                mermaid.append(f"        {data_id} --> {stage_id}")
+
+            # Update prev_output
+            prev_output = stage.get('output', f"Output{i}")
+
+        mermaid.append("    end")
+
+        result = "\n".join(mermaid)
+        logger.debug(f"Pipeline diagram generated with {len(stages)} stages")
+
+        return result
+
+    def extract_dataflow_from_parsed_files(self,
+                                             parsed_files: Dict[str, Dict],
+                                             max_entities: int = 30) -> Dict:
+        """
+        Extract data flow information from parsed code files.
+
+        Automatically identifies data entities and transformations.
+
+        Args:
+            parsed_files: Dictionary of file_path -> parsed code data
+            max_entities: Maximum number of entities to extract
+
+        Returns:
+            Dictionary with 'entities', 'transformations', and 'external_systems'
+        """
+        logger.info(f"Extracting dataflow from {len(parsed_files)} files")
+
+        data_entities = []
+        transformations = []
+        external_systems = []
+
+        # Track seen entities to avoid duplicates
+        seen_entities = set()
+        seen_transforms = set()
+
+        for file_path, parsed_data in parsed_files.items():
+            file_name = Path(file_path).stem
+
+            # Extract data entities from classes
+            classes = parsed_data.get('classes', [])
+
+            for cls in classes:
+                class_name = cls['name']
+
+                if class_name in seen_entities:
+                    continue
+
+                # Determine entity type based on naming patterns
+                entity_type = self._infer_entity_type(class_name, file_name)
+
+                data_entities.append({
+                    'name': class_name,
+                    'type': entity_type,
+                    'file': file_path
+                })
+
+                seen_entities.add(class_name)
+
+            # Extract transformations from functions
+            functions = parsed_data.get('functions', [])
+
+            for func in functions:
+                func_name = func['name']
+
+                # Skip private functions and test functions
+                if func_name.startswith('_') or func_name.startswith('test_'):
+                    continue
+
+                if func_name in seen_transforms:
+                    continue
+
+                # Try to infer inputs/outputs from function signature
+                inputs, outputs = self._infer_function_dataflow(func, seen_entities)
+
+                if inputs or outputs:
+                    transformations.append({
+                        'name': func_name,
+                        'inputs': inputs,
+                        'outputs': outputs,
+                        'file': file_path
+                    })
+
+                    seen_transforms.add(func_name)
+
+        # Limit to max_entities
+        if len(data_entities) > max_entities:
+            # Prioritize entities that are used in transformations
+            used_entities = set()
+            for t in transformations:
+                used_entities.update(t.get('inputs', []))
+                used_entities.update(t.get('outputs', []))
+
+            # Sort: used entities first, then by alphabetical
+            data_entities.sort(
+                key=lambda e: (e['name'] not in used_entities, e['name'])
+            )
+            data_entities = data_entities[:max_entities]
+
+        # Identify external systems (Database, API, Cache, etc.)
+        external_systems = self._identify_external_systems(parsed_files)
+
+        logger.debug(f"Extracted {len(data_entities)} entities, {len(transformations)} transformations")
+
+        return {
+            'entities': data_entities,
+            'transformations': transformations,
+            'external_systems': external_systems
+        }
+
+    # Helper methods
+
+    def _sanitize_id(self, name: str) -> str:
+        """Sanitize name for use as Mermaid ID."""
+        # Replace special characters with underscores
+        sanitized = re.sub(r'[^a-zA-Z0-9_]', '_', name)
+
+        # Ensure doesn't start with number
+        if sanitized and sanitized[0].isdigit():
+            sanitized = 'n_' + sanitized
+
+        return sanitized or 'unknown'
+
+    def _is_external_system(self,
+                             entity_name: str,
+                             external_systems: Optional[List[Dict]]) -> bool:
+        """Check if entity is an external system."""
+        if not external_systems:
+            return False
+
+        return any(sys['name'] == entity_name for sys in external_systems)
+
+    def _add_dataflow_styling(self,
+                                data_entities: List[Dict],
+                                transformations: List[Dict],
+                                external_systems: Optional[List[Dict]]) -> List[str]:
+        """Add styling for dataflow diagram."""
+        styles = []
+
+        # Define class styles
+        styles.append("    classDef dataEntity fill:#e1f5ff,stroke:#01579b,stroke-width:2px")
+        styles.append("    classDef transformation fill:#fff3e0,stroke:#e65100,stroke-width:2px")
+        styles.append("    classDef external fill:#f3e5f5,stroke:#4a148c,stroke-width:2px")
+
+        # Apply to data entities
+        entity_ids = [self._sanitize_id(e['name']) for e in data_entities]
+        if entity_ids:
+            styles.append(f"    class {','.join(entity_ids)} dataEntity")
+
+        # Apply to transformations
+        transform_ids = [self._sanitize_id(t['name']) for t in transformations]
+        if transform_ids:
+            styles.append(f"    class {','.join(transform_ids)} transformation")
+
+        # Apply to external systems
+        if external_systems:
+            external_ids = [self._sanitize_id(s['name']) for s in external_systems]
+            if external_ids:
+                styles.append(f"    class {','.join(external_ids)} external")
+
+        return styles
+
+    def _infer_entity_type(self, class_name: str, file_name: str) -> str:
+        """Infer entity type from class name and file name."""
+        class_lower = class_name.lower()
+        file_lower = file_name.lower()
+
+        # Check for common patterns
+        if 'dto' in class_lower or 'request' in class_lower or 'response' in class_lower:
+            return 'dto'
+        elif 'model' in class_lower or 'model' in file_lower or 'entity' in class_lower:
+            return 'model'
+        elif 'interface' in class_lower or 'protocol' in class_lower:
+            return 'interface'
+        elif 'config' in class_lower or 'settings' in class_lower:
+            return 'config'
+        else:
+            return 'data'
+
+    def _infer_function_dataflow(self,
+                                   func: Dict,
+                                   known_entities: Set[str]) -> tuple:
+        """
+        Infer function inputs and outputs from signature and name.
+
+        Returns:
+            Tuple of (inputs, outputs)
+        """
+        inputs = []
+        outputs = []
+
+        # Get parameters and return type if available
+        params = func.get('parameters', [])
+        return_type = func.get('return_type', '')
+
+        # Match parameters to known entities
+        for param in params:
+            param_name = param if isinstance(param, str) else param.get('name', '')
+            param_type = param.get('type', '') if isinstance(param, dict) else ''
+
+            # Check if parameter type matches a known entity
+            if param_type in known_entities:
+                inputs.append(param_type)
+            elif param_name.title() in known_entities:
+                inputs.append(param_name.title())
+
+        # Check if return type matches a known entity
+        if return_type in known_entities:
+            outputs.append(return_type)
+
+        # If no matches, try to infer from function name
+        if not inputs and not outputs:
+            func_name = func['name']
+
+            # Common transformation patterns
+            if func_name.startswith('create_') or func_name.startswith('build_'):
+                # Creates an entity
+                entity_name = func_name.replace('create_', '').replace('build_', '').title()
+                if entity_name in known_entities:
+                    outputs.append(entity_name)
+
+            elif func_name.startswith('process_') or func_name.startswith('transform_'):
+                # Processes an entity
+                entity_name = func_name.replace('process_', '').replace('transform_', '').title()
+                if entity_name in known_entities:
+                    inputs.append(entity_name)
+                    outputs.append(f"Processed{entity_name}")
+
+            elif '_to_' in func_name:
+                # Converts from one type to another
+                parts = func_name.split('_to_')
+                if len(parts) == 2:
+                    from_entity = parts[0].title()
+                    to_entity = parts[1].title()
+
+                    if from_entity in known_entities:
+                        inputs.append(from_entity)
+                    if to_entity in known_entities:
+                        outputs.append(to_entity)
+
+        return inputs, outputs
+
+    def _identify_external_systems(self, parsed_files: Dict[str, Dict]) -> List[Dict]:
+        """Identify external systems (databases, APIs, caches) from code."""
+        external_systems = []
+        seen_systems = set()
+
+        # Common patterns for external systems
+        db_keywords = ['database', 'db', 'session', 'connection', 'query']
+        api_keywords = ['api', 'client', 'request', 'endpoint']
+        cache_keywords = ['cache', 'redis', 'memcache']
+        queue_keywords = ['queue', 'broker', 'kafka', 'rabbitmq']
+
+        for file_path, parsed_data in parsed_files.items():
+            # Check imports for external systems
+            imports = parsed_data.get('imports', [])
+
+            for imp in imports:
+                imp_lower = imp.lower()
+
+                system_type = None
+                system_name = None
+
+                if any(kw in imp_lower for kw in db_keywords):
+                    system_type = 'storage'
+                    system_name = 'Database'
+                elif any(kw in imp_lower for kw in api_keywords):
+                    system_type = 'api'
+                    system_name = 'External API'
+                elif any(kw in imp_lower for kw in cache_keywords):
+                    system_type = 'cache'
+                    system_name = 'Cache'
+                elif any(kw in imp_lower for kw in queue_keywords):
+                    system_type = 'queue'
+                    system_name = 'Message Queue'
+
+                if system_name and system_name not in seen_systems:
+                    external_systems.append({
+                        'name': system_name,
+                        'type': system_type
+                    })
+                    seen_systems.add(system_name)
+
+        return external_systems