rust-crate-pipeline 1.1.0__py3-none-any.whl

rust_crate_pipeline/network.py

@@ -0,0 +1,307 @@
+# network.py
+import os
+import re
+import time
+import logging
+import requests
+from requests_cache import CachedSession
+from bs4 import BeautifulSoup
+from typing import Dict, List, Optional
+from .config import PipelineConfig
+
+class GitHubBatchClient:
+    def __init__(self, config: PipelineConfig):
+        self.config = config
+        self.headers = {"Accept": "application/vnd.github.v3+json"}
+        if config.github_token:
+            self.headers["Authorization"] = f"token {config.github_token}"
+        
+        self.session = CachedSession(
+            'github_cache', 
+            expire_after=config.cache_ttl * 2  # Longer cache for GitHub
+        )
+        self.remaining_calls = 5000
+        self.reset_time = 0
+
+    def check_rate_limit(self):
+        """Check and update current rate limit status"""
+        try:
+            response = self.session.get("https://api.github.com/rate_limit", headers=self.headers)
+            if response.ok:
+                data = response.json()
+                self.remaining_calls = data["resources"]["core"]["remaining"]
+                self.reset_time = data["resources"]["core"]["reset"]
+                
+                if self.remaining_calls < 100:
+                    reset_in = self.reset_time - time.time()
+                    logging.warning(f"GitHub API rate limit low: {self.remaining_calls} remaining. Resets in {reset_in/60:.1f} minutes")
+        except Exception:
+            pass
+
+    def get_repo_stats(self, owner: str, repo: str) -> Dict:
+        """Get repository statistics"""
+        try:
+            url = f"https://api.github.com/repos/{owner}/{repo}"
+            response = self.session.get(url, headers=self.headers)
+            if response.ok:
+                return response.json()
+            else:
+                logging.warning(f"Failed to get repo stats for {owner}/{repo}: {response.status_code}")
+                return {}
+        except Exception as e:
+            logging.error(f"Error fetching repo stats: {str(e)}")
+            return {}
+
+    def batch_get_repo_stats(self, repo_list: List[str]) -> Dict[str, Dict]:
+        """Get statistics for multiple repositories in a batch"""
+        self.check_rate_limit()
+        
+        results = {}
+        for repo_url in repo_list:
+            # Extract owner/repo from URL
+            match = re.search(r"github\.com/([^/]+)/([^/\.]+)", repo_url)
+            if not match:
+                continue
+                
+            owner, repo = match.groups()
+            repo = repo.split('.')[0]  # Remove .git extension if present
+            
+            # Get stats
+            stats = self.get_repo_stats(owner, repo)
+            results[repo_url] = stats
+            
+            # Be nice to GitHub API
+            time.sleep(0.1)
+        
+        return results
+
+class CrateAPIClient:
+    def __init__(self, config: PipelineConfig):
+        self.config = config
+        self.session = CachedSession('crate_cache', expire_after=config.cache_ttl)
+        
+    def fetch_crate_metadata(self, crate_name: str) -> Optional[Dict]:
+        """Fetch metadata with retry logic"""
+        for attempt in range(self.config.max_retries):
+            try:
+                return self._fetch_metadata(crate_name)
+            except Exception as e:
+                logging.warning(f"Attempt {attempt+1} failed for {crate_name}: {str(e)}")
+                wait = 2 ** attempt
+                time.sleep(wait)
+        return None
+
+    def _fetch_metadata(self, crate_name: str) -> Optional[Dict]:
+        """Enhanced metadata fetching that tries multiple sources"""
+        # First try crates.io (primary source)
+        try:
+            r = self.session.get(f"https://crates.io/api/v1/crates/{crate_name}")
+            if r.ok:
+                data = r.json()
+                crate_data = data["crate"]
+                latest = crate_data["newest_version"]
+                
+                # Get readme
+                readme_response = self.session.get(f"https://crates.io/api/v1/crates/{crate_name}/readme")
+                readme = readme_response.text if readme_response.ok else ""
+                
+                # Get dependencies
+                deps_response = self.session.get(f"https://crates.io/api/v1/crates/{crate_name}/{latest}/dependencies")
+                deps = deps_response.json().get("dependencies", []) if deps_response.ok else []
+                
+                # Get features - using the versions endpoint
+                features = []
+                versions_response = self.session.get(f"https://crates.io/api/v1/crates/{crate_name}/{latest}")
+                if versions_response.ok:
+                    version_data = versions_response.json().get("version", {})
+                    features_dict = version_data.get("features", {})
+                    features = [{"name": k, "dependencies": v} for k, v in features_dict.items()]
+                
+                # Repository info and GitHub stars
+                repo = crate_data.get("repository", "")
+                gh_stars = 0
+                
+                # Check if it's a GitHub repo
+                if "github.com" in repo and self.config.github_token:
+                    match = re.search(r"github.com/([^/]+)/([^/]+)", repo)
+                    if match:
+                        owner, repo_name = match.groups()
+                        repo_name = repo_name.split('.')[0]  # Handle .git extensions
+                        gh_url = f"https://api.github.com/repos/{owner}/{repo_name}"
+                        gh_headers = {"Authorization": f"token {self.config.github_token}"} if self.config.github_token else {}
+                        gh = self.session.get(gh_url, headers=gh_headers)
+                        if gh.ok:
+                            gh_data = gh.json()
+                            gh_stars = gh_data.get("stargazers_count", 0)
+                
+                # Check if it's hosted on lib.rs
+                lib_rs_data = {}
+                if "lib.rs" in repo:
+                    lib_rs_url = f"https://lib.rs/crates/{crate_name}"
+                    lib_rs_response = self.session.get(lib_rs_url)
+                    if lib_rs_response.ok:
+                        soup = BeautifulSoup(lib_rs_response.text, 'html.parser')
+                        # Get README from lib.rs if not already available
+                        if not readme:
+                            readme_div = soup.find('div', class_='readme')
+                            if readme_div:
+                                readme = readme_div.get_text(strip=True)
+                        
+                        # Get lib.rs specific stats
+                        stats_div = soup.find('div', class_='crate-stats')
+                        if stats_div:
+                            downloads_text = stats_div.find(string=re.compile(r'[\d,]+ downloads'))
+                            if downloads_text:
+                                lib_rs_data["librs_downloads"] = int(re.sub(r'[^\d]', '', downloads_text))
+                
+                # Extract code snippets from readme
+                code_snippets = self.extract_code_snippets(readme)
+                
+                # Extract sections from readme
+                readme_sections = self.extract_readme_sections(readme) if readme else {}
+                
+                result = {
+                    "name": crate_name,
+                    "version": latest,
+                    "description": crate_data.get("description", ""),
+                    "repository": repo,
+                    "keywords": crate_data.get("keywords", []),
+                    "categories": crate_data.get("categories", []),
+                    "readme": readme,
+                    "downloads": crate_data.get("downloads", 0),
+                    "github_stars": gh_stars,
+                    "dependencies": deps,
+                    "code_snippets": code_snippets,
+                    "features": features,
+                    "readme_sections": readme_sections,
+                    **lib_rs_data
+                }
+                
+                return result
+                
+        except Exception as e:
+            logging.error(f"Failed fetching metadata for {crate_name}: {str(e)}")
+            raise
+        
+        # If crates.io fails, try lib.rs
+        try:
+            r = self.session.get(f"https://lib.rs/crates/{crate_name}")
+            if r.ok:
+                soup = BeautifulSoup(r.text, 'html.parser')
+                
+                # Extract metadata from lib.rs page
+                name = soup.select_one('h1').text.strip() if soup.select_one('h1') else crate_name
+                
+                # Find description
+                desc_elem = soup.select_one('.description')
+                description = desc_elem.text.strip() if desc_elem else ""
+                
+                # Find repository link
+                repo_link = None
+                for a in soup.select('a'):
+                    if 'github.com' in a.get('href', ''):
+                        repo_link = a['href']
+                        break
+                
+                # Basic metadata from lib.rs
+                return {
+                    "name": name,
+                    "version": "latest",  # lib.rs doesn't easily expose version
+                    "description": description,
+                    "repository": repo_link or "",
+                    "keywords": [],
+                    "categories": [],
+                    "readme": "",
+                    "downloads": 0,
+                    "github_stars": 0,
+                    "dependencies": [],
+                    "code_snippets": [],
+                    "features": [],
+                    "readme_sections": {},
+                    "source": "lib.rs",
+                }
+        except Exception:
+            pass
+        
+        # Finally, try GitHub search
+        try:
+            # This is a simplification - GitHub's search API requires authentication
+            headers = {}
+            if self.config.github_token:
+                headers["Authorization"] = f"token {self.config.github_token}"
+                
+            search_url = f"https://api.github.com/search/repositories?q={crate_name}+language:rust"
+            r = requests.get(search_url, headers=headers)
+            
+            if r.ok:
+                results = r.json().get("items", [])
+                if results:
+                    repo = results[0]  # Take first match
+                    
+                    # Basic metadata from GitHub
+                    return {
+                        "name": crate_name,
+                        "version": "unknown",
+                        "description": repo.get("description", ""),
+                        "repository": repo.get("html_url", ""),
+                        "keywords": [],
+                        "categories": [],
+                        "readme": "",
+                        "downloads": 0,
+                        "github_stars": repo.get("stargazers_count", 0),
+                        "dependencies": [],
+                        "code_snippets": [],
+                        "features": [],
+                        "readme_sections": {},
+                        "source": "github",
+                    }
+        except Exception:
+            pass
+        
+        # If all sources fail
+        return None
+
+    def extract_code_snippets(self, readme: str) -> List[str]:
+        """Extract code snippets from markdown README"""
+        snippets = []
+        if not readme:
+            return snippets
+            
+        # Find Rust code blocks
+        pattern = r"```(?:rust|(?:no_run|ignore|compile_fail|mdbook-runnable)?)\s*([\s\S]*?)```"
+        matches = re.findall(pattern, readme)
+        
+        for code in matches:
+            if len(code.strip()) > 10:  # Only include non-trivial snippets
+                snippets.append(code.strip())
+        
+        return snippets[:5]  # Limit to 5 snippets
+
+    def extract_readme_sections(self, readme: str) -> Dict[str, str]:
+        """Extract sections from README based on markdown headers"""
+        if not readme:
+            return {}
+            
+        sections = {}
+        lines = readme.split('\n')
+        current_section = ""
+        current_content = []
+        
+        for line in lines:
+            if re.match(r'^#+\s+', line):  # It's a header
+                # Save previous section
+                if current_section and current_content:
+                    sections[current_section] = '\n'.join(current_content).strip()
+                
+                # Start new section
+                current_section = re.sub(r'^#+\s+', '', line).strip()
+                current_content = []
+            else:
+                if current_section:  # Only collect content if we have a section
+                    current_content.append(line)
+        
+        # Don't forget the last section
+        if current_section and current_content:
+            sections[current_section] = '\n'.join(current_content).strip()
+        
+        return sections

rust_crate_pipeline/pipeline.py

@@ -0,0 +1,260 @@
+# pipeline.py
+import os
+import time
+import logging
+import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+from typing import List, Dict, Optional
+from .config import PipelineConfig, CrateMetadata, EnrichedCrate
+from .network import CrateAPIClient, GitHubBatchClient
+from .ai_processing import LLMEnricher
+from .analysis import SourceAnalyzer, SecurityAnalyzer, UserBehaviorAnalyzer, DependencyAnalyzer
+
+class CrateDataPipeline:
+    def __init__(self, config: PipelineConfig):
+        self.config = config
+        self.api_client = CrateAPIClient(config)
+        self.github_client = GitHubBatchClient(config)
+        self.enricher = LLMEnricher(config)
+        self.crates = self.get_crate_list()
+        self.output_dir = self._create_output_dir()
+        
+    def _create_output_dir(self) -> str:
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        output_dir = f"crate_data_{timestamp}"
+        os.makedirs(output_dir, exist_ok=True)
+        return output_dir
+
+    def get_crate_list(self, limit: Optional[int] = None) -> List[str]:
+        """Return a comprehensive list of all high-value crates to process"""
+        crates = [
+            # Web frameworks
+            "actix-web", "rocket", "axum", "warp", "tower",
+            
+            # Async runtimes and utilities
+            "tokio", "tokio-stream", "async-trait", "futures",
+            
+            # Serialization/deserialization
+            "serde", "serde_json", "serde_yaml", "bincode",
+            
+            # Error handling
+            "anyhow", "thiserror",
+            
+            # Utilities
+            "rand", "uuid", "chrono", "regex", "log", "env_logger", "clap", "crossterm", 
+            "itertools", "num", "cfg-if", "bytes", "mime", "form_urlencoded", "parking_lot", 
+            "csv", "lazy_static", "once_cell", "tracing", "base64", "sha2", "flate2", "tar",
+            
+            # HTTP clients and servers
+            "reqwest", "hyper",
+            
+            # Database
+            "sqlx", "diesel", "postgres", "rusqlite",
+            
+            # Concurrency
+            "rayon",
+            
+            # Protocol buffers and gRPC
+            "prost", "tonic",
+            
+            # Procedural macros
+            "syn", "quote", "proc-macro2",
+            
+            # Machine Learning & AI
+            "tokenizers", "safetensors", "linfa", "ndarray", "smartcore", "burn", 
+            "tract-core", "tract-onnx", "tract-hir", "tract-linalg", "tract-data", 
+            "tract-nnef", "tract-onnx-opl", "tract-pulse", "tract-pulse-opl", 
+            "tract-nnef-resources", "tch", "torch-sys", "ort", "ort-sys", "candle-core", 
+            "candle-nn", "candle-transformers", "candle-kernels", "candle-onnx", 
+            "candle-metal-kernels", "tiktoken-rs", "tensorflow", "tensorflow-sys",
+            "onnxruntime", "onnxruntime-sys", "onnx-protobuf", "llama-cpp-2", 
+            "llama-cpp-sys-2", "llm", "llm-samplers", "llm-chain", "llm-chain-openai",
+            "llama-core", "llamaedge", "openai", "openai-api-rs", "openai_dive", 
+            "genai", "aleph-alpha-client", "llm_api_access", "ollama-rs", 
+            "rust-bert", "fastembed", "hf-hub", "whisper-rs-sys", "toktrie", 
+            "toktrie_hf_tokenizers", "toktrie_hf_downloader", "rust_tokenizers",
+        ]
+        
+        if limit is not None:
+            return crates[:limit]
+        return crates
+
+    def fetch_metadata_batch(self, crate_names: List[str]) -> List[CrateMetadata]:
+        """Fetch metadata for a batch of crates in parallel"""
+        with ThreadPoolExecutor(max_workers=self.config.n_workers) as executor:
+            futures = {executor.submit(self.api_client.fetch_crate_metadata, name): name 
+                      for name in crate_names}
+            
+            results = []
+            for future in as_completed(futures):
+                crate_name = futures[future]
+                try:
+                    data = future.result()
+                    if data:
+                        # Convert dict to CrateMetadata
+                        crate_metadata = CrateMetadata(
+                            name=data.get("name", ""),
+                            version=data.get("version", ""),
+                            description=data.get("description", ""),
+                            repository=data.get("repository", ""),
+                            keywords=data.get("keywords", []),
+                            categories=data.get("categories", []),
+                            readme=data.get("readme", ""),
+                            downloads=data.get("downloads", 0),
+                            github_stars=data.get("github_stars", 0),
+                            dependencies=data.get("dependencies", []),
+                            features=data.get("features", []),
+                            code_snippets=data.get("code_snippets", []),
+                            readme_sections=data.get("readme_sections", {}),
+                            librs_downloads=data.get("librs_downloads"),
+                            source=data.get("source", "crates.io")
+                        )
+                        results.append(crate_metadata)
+                        logging.info(f"Fetched metadata for {crate_name}")
+                except Exception as e:
+                    logging.error(f"Failed to fetch metadata for {crate_name}: {str(e)}")
+            
+            return results
+
+    def enrich_batch(self, batch: List[CrateMetadata]) -> List[EnrichedCrate]:
+        """Enrich a batch of crates with GitHub stats and AI"""
+        # Add GitHub stats first
+        github_repos = [c.repository for c in batch if "github.com" in c.repository]
+        repo_stats = self.github_client.batch_get_repo_stats(github_repos)
+        
+        # Update crates with GitHub info
+        for crate in batch:
+            repo_url = crate.repository
+            if repo_url in repo_stats:
+                stats = repo_stats[repo_url]
+                crate.github_stars = stats.get("stargazers_count", 0)
+        
+        # Now enrich with AI
+        enriched_batch = []
+        for crate in batch:
+            try:
+                enriched = self.enricher.enrich_crate(crate)
+                enriched_batch.append(enriched)
+                logging.info(f"Enriched {crate.name}")
+            except Exception as e:
+                logging.error(f"Failed to enrich {crate.name}: {str(e)}")
+                # Add the crate with just the fields we have
+                enriched_dict = crate.__dict__.copy()
+                enriched_batch.append(EnrichedCrate(**enriched_dict))
+                
+        return enriched_batch
+
+    def analyze_dependencies(self, crates: List[EnrichedCrate]) -> Dict:
+        """Analyze dependencies between crates"""
+        return DependencyAnalyzer.analyze_dependencies(crates)
+
+    def save_checkpoint(self, data: List[EnrichedCrate], prefix: str):
+        """Save processing checkpoint with status metadata"""
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        filename = os.path.join(self.output_dir, f"{prefix}_{timestamp}.jsonl")
+        
+        with open(filename, "w") as f:
+            for item in data:
+                # Convert to dict for serialization
+                item_dict = item.__dict__.copy()
+                f.write(json.dumps(item_dict) + "\n")
+        
+        # Save status metadata
+        status = {
+            "timestamp": timestamp,
+            "total_crates": len(data),
+            "processed_crates": sum(1 for c in data if c.use_case is not None),
+            "advanced_analysis": sum(1 for c in data if c.source_analysis is not None),
+            "checkpoint_file": filename
+        }
+        
+        status_file = os.path.join(self.output_dir, f"{prefix}_status_{timestamp}.json")
+        with open(status_file, "w") as f:
+            json.dump(status, f, indent=2)
+            
+        logging.info(f"Saved checkpoint to {filename}")
+        return filename
+
+    def save_final_output(self, data: List[EnrichedCrate], dependency_data: Dict):
+        """Save final enriched data and analysis"""
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        
+        # Save main enriched data
+        final_output = os.path.join(self.output_dir, f"enriched_crate_metadata_{timestamp}.jsonl")
+        with open(final_output, "w") as f:
+            for item in data:
+                item_dict = item.__dict__.copy()
+                f.write(json.dumps(item_dict) + "\n")
+                
+        # Save dependency analysis
+        dep_file = os.path.join(self.output_dir, f"dependency_analysis_{timestamp}.json")
+        with open(dep_file, "w") as f:
+            json.dump(dependency_data, f, indent=2)
+
+        # Generate summary report
+        summary = {
+            "total_crates": len(data),
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+            "most_popular": sorted([{
+                "name": c.name,
+                "score": c.score or 0,
+                "downloads": c.downloads,
+                "github_stars": c.github_stars
+            } for c in data], key=lambda x: x["score"], reverse=True)[:5],
+            "most_depended_upon": dependency_data.get("most_depended", [])[:5]
+        }
+        
+        summary_file = os.path.join(self.output_dir, f"summary_report_{timestamp}.json")
+        with open(summary_file, "w") as f:
+            json.dump(summary, f, indent=2)
+
+        logging.info(f"Results saved to {self.output_dir}/")
+
+    def run(self):
+        """Main pipeline execution flow"""
+        start_time = time.time()
+        logging.info(f"Processing {len(self.crates)} crates...")
+        
+        # Process in batches
+        all_enriched = []
+        crate_batches = [self.crates[i:i+self.config.batch_size] 
+                         for i in range(0, len(self.crates), self.config.batch_size)]
+
+        for batch_num, batch in enumerate(crate_batches):
+            logging.info(f"Processing batch {batch_num+1}/{len(crate_batches)} ({len(batch)} crates)")
+            
+            # Fetch metadata
+            batch_data = self.fetch_metadata_batch(batch)
+            
+            # Enrich the batch
+            enriched_batch = self.enrich_batch(batch_data)
+            all_enriched.extend(enriched_batch)
+            
+            # Save checkpoint after each batch
+            self.save_checkpoint(all_enriched, "batch_checkpoint")
+            logging.info(f"Completed batch {batch_num+1}, processed {len(all_enriched)}/{len(self.crates)} crates so far")
+            
+            # Optional: Add source analysis for some crates
+            if batch_num < 2:  # Only do detailed analysis for first 2 batches
+                for crate in enriched_batch:
+                    try:
+                        crate.source_analysis = SourceAnalyzer.analyze_crate_source(crate)
+                        crate.security = SecurityAnalyzer.check_security_metrics(crate)
+                        crate.user_behavior = UserBehaviorAnalyzer.fetch_user_behavior_data(crate)
+                        logging.info(f"Advanced analysis completed for {crate.name}")
+                    except Exception as e:
+                        logging.warning(f"Advanced analysis failed for {crate.name}: {str(e)}")
+        
+        # Step 3: Perform dependency analysis
+        logging.info("Analyzing crate dependencies...")
+        dependency_analysis = self.analyze_dependencies(all_enriched)
+        
+        # Save final results
+        self.save_final_output(all_enriched, dependency_analysis)
+
+        # Final summary
+        duration = time.time() - start_time
+        logging.info(f"✅ Done. Enriched {len(all_enriched)} crates in {duration:.2f}s")
+        
+        return all_enriched, dependency_analysis

rust_crate_pipeline/utils/file_utils.py

@@ -0,0 +1,72 @@
+# rust_crate_pipeline/utils/file_utils.py
+import os
+import json
+import shutil
+from datetime import datetime
+from typing import List, Dict, Any
+
+def create_output_dir(base_name: str = "crate_data") -> str:
+    """
+    Create timestamped output directory
+    
+    Args:
+        base_name: Base name for output directory
+    
+    Returns:
+        Path to created directory
+    """
+    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+    output_dir = f"{base_name}_{timestamp}"
+    os.makedirs(output_dir, exist_ok=True)
+    return output_dir
+
+def save_checkpoint(data: List[Dict], prefix: str, output_dir: str) -> str:
+    """
+    Save processing checkpoint with status metadata
+    
+    Args:
+        data: List of crate dictionaries
+        prefix: File name prefix
+        output_dir: Target directory
+    
+    Returns:
+        Path to saved checkpoint file
+    """
+    timestamp = datetime.now().isoformat()
+    filename = os.path.join(output_dir, f"{prefix}_{timestamp}.jsonl")
+    
+    with open(filename, "w") as f:
+        for item in data:
+            f.write(json.dumps(item) + "\n")
+    
+    # Save status metadata
+    status = {
+        "timestamp": timestamp,
+        "total_items": len(data),
+        "checkpoint_file": filename
+    }
+    
+    status_file = os.path.join(output_dir, f"{prefix}_status_{timestamp}.json")
+    with open(status_file, "w") as f:
+        json.dump(status, f, indent=2)
+    
+    return filename
+
+def safe_file_cleanup(path: str):
+    """Safely remove files or directories"""
+    try:
+        if os.path.isfile(path):
+            os.remove(path)
+        elif os.path.isdir(path):
+            shutil.rmtree(path)
+    except Exception as e:
+        print(f"Failed to cleanup {path}: {str(e)}")
+
+def disk_space_check(min_free_gb: float = 1.0) -> bool:
+    """Check if sufficient disk space is available"""
+    try:
+        free_bytes = shutil.disk_usage(".").free
+        free_gb = free_bytes / (1024 ** 3)
+        return free_gb >= min_free_gb
+    except Exception:
+        return True  # Assume OK if check fails