diff --git a/.gitignore b/.gitignore index fbc8025..af92ec7 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,7 @@ vendor/ .DS_Store Thumbs.db data/ + +bootstrap +*.run +*.tar.gz diff --git a/archive.py b/archive.py index 2bfa68f..3b8e6e7 100644 --- a/archive.py +++ b/archive.py @@ -43,7 +43,7 @@ import aiofiles.os from bs4 import BeautifulSoup from tqdm import tqdm -from neopig import NeoPig, setup_logging +from neopig import NeoPig, setup_logging, rotate_state_file, get_state_file_path from async_web_fetcher import CrawlMode from screenshot import ScreenshotConfig @@ -111,6 +111,8 @@ class SiteArchiver: screenshot_config: ScreenshotConfig = None, fast_mode: bool = False, trim_wrapper: bool = False, + show_progress: bool = True, + fresh_start: bool = False, ): self.output_dir = Path(output_dir) self.include_screenshots = include_screenshots @@ -118,6 +120,8 @@ class SiteArchiver: self.screenshot_config = screenshot_config or ScreenshotConfig(enabled=include_screenshots) self.fast_mode = fast_mode self.trim_wrapper = trim_wrapper + self.show_progress = show_progress + self.fresh_start = fresh_start async def archive( self, @@ -157,14 +161,19 @@ class SiteArchiver: ) await pig.init() - # Load previously crawled media/screenshots from DB (source of truth for resume) - crawled_media = await pig.db.get_crawled_media_uris() - crawled_screenshots = await pig.db.get_crawled_screenshot_uris() - logger.info(f"Resume state: {len(crawled_media)} media URIs, {len(crawled_screenshots)} screenshotted pages in DB") - if crawled_media: - pig.seen_media = crawled_media - if crawled_screenshots: - pig.seen_screenshots = crawled_screenshots + # Handle fresh start: rotate state files + if self.fresh_start: + pig._clear_state(target_url) + logger.info("Fresh start: state file rotated") + else: + # Load previously crawled media/screenshots from DB (source of truth for resume) + crawled_media = await pig.db.get_crawled_media_uris() + crawled_screenshots = await pig.db.get_crawled_screenshot_uris() + logger.info(f"Resume state: {len(crawled_media)} media URIs, {len(crawled_screenshots)} screenshotted pages in DB") + if crawled_media: + pig.seen_media = crawled_media + if crawled_screenshots: + pig.seen_screenshots = crawled_screenshots logger.info(f"Screenshots enabled: {pig.screenshot_config.enabled}") # Run the crawl @@ -247,8 +256,7 @@ class SiteArchiver: (tmpdir_path / 'metadata.json').write_text(json.dumps(metadata, indent=2)) # Copy state file - state_domain = domain.replace('.', '-').replace(':', '-') - state_file = Path("data") / f"crawl-state-{state_domain}.json" + state_file = get_state_file_path(domain) if state_file.exists(): shutil.copy(state_file, tmpdir_path / 'crawl_state.json') logger.info(f"Included crawl state for future delta crawls") @@ -259,6 +267,17 @@ class SiteArchiver: # Step 3: Stream everything to tar.gz in one pass logger.info("Streaming to archive...") + # Capture for closure + show_progress = self.show_progress + include_markdown = self.include_markdown + trim_wrapper = self.trim_wrapper + include_screenshots = self.include_screenshots + + # Pre-count screenshots for progress bar + screenshot_files = [] + if include_screenshots and linkpeek_vault_path.exists(): + screenshot_files = list(linkpeek_vault_path.rglob('*.png')) + list(linkpeek_vault_path.rglob('*.jpg')) + def stream_to_tar(): with tarfile.open(tar_path, 'w:gz') as tar: # Add generated files first (from temp) @@ -266,7 +285,8 @@ class SiteArchiver: tar.add(f, arcname=f"{archive_name}/{f.name}") # Stream HTML files (with rewritten URLs) - for rel_path_str, content in html_contents.items(): + html_iter = tqdm(html_contents.items(), desc="HTML", unit="pages", disable=not show_progress) + for rel_path_str, content in html_iter: try: html_bytes = content.encode('utf-8') arcname = f"{archive_name}/html/{rel_path_str}" @@ -275,8 +295,8 @@ class SiteArchiver: tar.addfile(info, io.BytesIO(html_bytes)) # Generate markdown on the fly - if self.include_markdown: - md_content = html_to_markdown(content, trim_wrapper=self.trim_wrapper) + if include_markdown: + md_content = html_to_markdown(content, trim_wrapper=trim_wrapper) md_bytes = md_content.encode('utf-8') md_rel = rel_path_str.replace('.html', '.md') md_arcname = f"{archive_name}/markdown/{md_rel}" @@ -287,10 +307,10 @@ class SiteArchiver: logger.debug(f"Error adding HTML {rel_path_str}: {e}") # Stream media files by hash (matching rewritten URLs) - # Build hash -> file path mapping from hash vault hash_vault = Path(vault_path) added_hashes = set() - for url, md5 in url_to_hash.items(): + media_iter = tqdm(url_to_hash.items(), desc="Media", unit="files", disable=not show_progress) + for url, md5 in media_iter: if md5 in added_hashes: continue # Find file in hash vault: vault/xx/hash.ext @@ -308,8 +328,9 @@ class SiteArchiver: break # Stream screenshots (follow symlinks) - if self.include_screenshots and linkpeek_vault_path.exists(): - for screenshot_file in linkpeek_vault_path.rglob('*.png'): + if include_screenshots and screenshot_files: + screen_iter = tqdm(screenshot_files, desc="Screenshots", unit="files", disable=not show_progress) + for screenshot_file in screen_iter: try: rel_path = screenshot_file.relative_to(linkpeek_vault_path) arcname = f"{archive_name}/screenshots/{rel_path}" @@ -526,74 +547,116 @@ class SiteArchiver: # See: make run TARBALL=archive.tar.gz -def upgrade_neopig_in_archive(archive_path: Path) -> Path: - """Replace neopig/ directory in existing archive with current source.""" +def upgrade_neopig_in_archive(archive_path: Path, output_dir: Path = None) -> Path: + """Replace neopig/ directory in existing archive with current source. + + Uses system tar + pigz for speed (10x faster than Python tarfile). + Requires temp disk space for extraction (~5-10x compressed size). + + Args: + archive_path: Path to the tar.gz archive + output_dir: Directory for temp files (default: cwd) + """ import shutil + import subprocess if not archive_path.exists(): raise FileNotFoundError(f"Archive not found: {archive_path}") + # Estimate required temp space (compressed * 5 is conservative) + archive_size = archive_path.stat().st_size + required_space = archive_size * 5 + logger.info(f"Archive: {archive_size / 1024 / 1024:.0f} MB (need ~{required_space / 1024 / 1024 / 1024:.1f} GB temp space)") + neopig_src = Path(__file__).parent - neopig_files = [ - 'neopig.py', 'database.py', 'async_web_fetcher.py', - 'storage.py', 'domain_vault.py', 'screenshot.py', - 'html2md.py', 'serp.py', 'filevault.py', 'async_filevault.py', - ] - # Create new archive with updated neopig - output_path = archive_path.with_suffix('.upgraded.tar.gz') + # Check for pigz (parallel gzip) - much faster + has_pigz = shutil.which('pigz') is not None + if has_pigz: + logger.info("Using pigz for parallel compression") - logger.info(f"Reading archive: {archive_path}") - with tarfile.open(archive_path, 'r:gz') as old_tar: - # Get members list (this reads the whole index) - members = old_tar.getmembers() - total = len(members) - archive_name = members[0].name.split('/')[0] - logger.info(f"Found {total} members in {archive_name}") + # Create temp directory for extraction (use output_dir or cwd for space) + work_dir = Path(output_dir) if output_dir else Path.cwd() + temp_base = work_dir / '.upgrade_tmp' + temp_base.mkdir(exist_ok=True) + with tempfile.TemporaryDirectory(dir=temp_base) as tmpdir: + tmpdir = Path(tmpdir) - with tarfile.open(output_path, 'w:gz') as new_tar: - # Copy all members except neopig/ and requirements.txt - copied = 0 - skipped = 0 - for i, member in enumerate(members): - parts = member.name.split('/') - if len(parts) > 1 and parts[1] == 'neopig': - skipped += 1 - continue # Skip old neopig files - if len(parts) > 1 and parts[1] == 'requirements.txt': - skipped += 1 - continue # Skip old requirements - if member.isfile(): - f = old_tar.extractfile(member) - if f: - new_tar.addfile(member, f) - copied += 1 - else: - new_tar.addfile(member) - copied += 1 + # Extract archive using system tar (much faster than Python) + logger.info(f"Extracting archive: {archive_path}") + subprocess.run(['tar', '-xzf', str(archive_path), '-C', str(tmpdir)], check=True) - if (i + 1) % 1000 == 0: - logger.info(f" Progress: {i+1}/{total} ({copied} copied, {skipped} skipped)") + # Find archive root directory + contents = list(tmpdir.iterdir()) + if len(contents) != 1 or not contents[0].is_dir(): + raise ValueError("Expected single directory in archive") + archive_root = contents[0] + archive_name = archive_root.name + logger.info(f"Archive root: {archive_name}") - logger.info(f"Copied {copied} members, skipped {skipped}") + # Remove old neopig directory + old_neopig = archive_root / 'neopig' + if old_neopig.exists(): + shutil.rmtree(old_neopig) + logger.info("Removed old neopig/") - # Add new neopig files - logger.info("Adding neopig source files...") - for pyfile in neopig_files: - src_path = neopig_src / pyfile - if src_path.exists(): - new_tar.add(src_path, arcname=f"{archive_name}/neopig/{pyfile}") - logger.info(f" Added: neopig/{pyfile}") + # Copy new neopig files + new_neopig = archive_root / 'neopig' + new_neopig.mkdir() + neopig_files = [ + 'neopig.py', 'database.py', 'async_web_fetcher.py', + 'storage.py', 'domain_vault.py', 'screenshot.py', + 'html2md.py', 'serp.py', 'filevault.py', 'async_filevault.py', + ] + for pyfile in neopig_files: + src = neopig_src / pyfile + if src.exists(): + shutil.copy2(src, new_neopig / pyfile) + logger.info(f" Added: neopig/{pyfile}") - # Add requirements.txt - req_path = neopig_src / 'requirements.txt' - if req_path.exists(): - new_tar.add(req_path, arcname=f"{archive_name}/requirements.txt") - logger.info(" Added: requirements.txt") + # Copy requirements.txt + req_src = neopig_src / 'requirements.txt' + if req_src.exists(): + shutil.copy2(req_src, archive_root / 'requirements.txt') + logger.info(" Added: requirements.txt") - # Replace original with upgraded - logger.info(f"Replacing original archive...") - shutil.move(output_path, archive_path) + # Repack using system tar (with pigz if available) + logger.info("Repacking archive...") + output_path = archive_path.with_suffix('.new.tar.gz') + + if has_pigz: + # tar + pigz for parallel compression + with open(output_path, 'wb') as out: + tar_proc = subprocess.Popen( + ['tar', '-cf', '-', '-C', str(tmpdir), archive_name], + stdout=subprocess.PIPE + ) + pigz_proc = subprocess.Popen( + ['pigz', '-c'], + stdin=tar_proc.stdout, + stdout=out + ) + tar_proc.stdout.close() + pigz_proc.wait() + tar_proc.wait() + else: + # Standard tar with gzip + subprocess.run( + ['tar', '-czf', str(output_path), '-C', str(tmpdir), archive_name], + check=True + ) + + # Replace original + logger.info("Replacing original archive...") + shutil.move(output_path, archive_path) + + # Cleanup temp base + try: + temp_base.rmdir() + except OSError: + pass + + logger.info(f"Done! Upgraded: {archive_path}") return archive_path @@ -613,7 +676,7 @@ async def main(): parser.add_argument("--screenshot-height", type=int, default=1024, help="Screenshot height") parser.add_argument("--screenshot-engine", type=str, default=None, help="Screenshot engine") parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output") - parser.add_argument("--fresh", action="store_true", help="Start fresh, ignore resume state") + parser.add_argument("--fresh", action="store_true", help="Start fresh (rotates old state files instead of resuming)") parser.add_argument("--fast", action="store_true", help="Fast mode: no crawl delay (for sites without robots.txt)") parser.add_argument("--serve", action="store_true", help="Start SERP server to watch crawl live") parser.add_argument("--port", type=int, default=8000, help="Port for SERP server (default: 8000)") @@ -621,6 +684,7 @@ async def main(): parser.add_argument("--trim-wrapper", action="store_true", help="With --backfill-markdown: strip nav/header/footer/logo before conversion") parser.add_argument("--db", default="data/neopig.db", help="Database path (for --backfill-markdown)") parser.add_argument("--upgrade-neopig", metavar="TARBALL", help="Upgrade neopig inside an existing archive") + parser.add_argument("--no-progress", action="store_true", help="Disable progress bars") args = parser.parse_args() @@ -629,8 +693,9 @@ async def main(): # Handle --upgrade-neopig if args.upgrade_neopig: archive_path = Path(args.upgrade_neopig) + output_dir = Path(args.output) if args.output != "." else None logger.info(f"Upgrading neopig in {archive_path}...") - upgrade_neopig_in_archive(archive_path) + upgrade_neopig_in_archive(archive_path, output_dir=output_dir) logger.info(f"Done! Archive updated: {archive_path}") return @@ -664,6 +729,8 @@ async def main(): screenshot_config=screenshot_config, fast_mode=args.fast, trim_wrapper=args.trim_wrapper, + show_progress=not args.no_progress, + fresh_start=args.fresh, ) # Start SERP server if requested diff --git a/html2md.py b/html2md.py index 4f53249..816ae0f 100644 --- a/html2md.py +++ b/html2md.py @@ -415,7 +415,8 @@ class SmartMarkdownConverter: src = self._resolve_url(child.get('src', '')) alt = child.get('alt', '') if src: - parts.append(f"![{alt}]({src})") + # Images should be block-level, not inline + parts.append(f"\n\n![{alt}]({src})\n\n") elif tag == 'br': parts.append('\n') diff --git a/neopig.py b/neopig.py index 46b4df9..36cc8d8 100644 --- a/neopig.py +++ b/neopig.py @@ -25,7 +25,7 @@ import json import logging import os import sys -from datetime import datetime, timezone +from datetime import datetime, timezone, timezone from pathlib import Path from typing import List, Dict, Any, Optional, Set @@ -48,6 +48,56 @@ from tqdm import tqdm logger = logging.getLogger(__name__) +def get_state_file_path(domain: str) -> Path: + """Get unified state file path for a domain. + + Args: + domain: Domain name (e.g., 'example.com') + + Returns: + Path like data/{domain}.state + """ + safe_domain = domain.replace('://', '-').replace('/', '-').replace('.', '-') + return Path(f"data/{safe_domain}.state") + + +def rotate_state_file(state_path: Path, preserve_keys: List[str] = None) -> Optional[Path]: + """Rotate state file with optional selective preservation. + + Args: + state_path: Path to state file + preserve_keys: If provided, rotate then copy back these keys from rotated file. + If None, just rotate (full fresh start). + + Returns: + Path to rotated file, or None if no rotation needed. + """ + if not state_path.exists(): + return None + + # Find next available rotation number + i = 1 + while Path(f"{state_path}.{i}").exists(): + i += 1 + rotated = Path(f"{state_path}.{i}") + state_path.rename(rotated) + logger.info(f"Rotated state file to {rotated}") + + # If preserve_keys specified, copy back those keys from rotated file + if preserve_keys: + try: + import json + old_state = json.loads(rotated.read_text()) + new_state = {k: v for k, v in old_state.items() if k in preserve_keys} + if new_state: + state_path.write_text(json.dumps(new_state, indent=2)) + logger.info(f"Preserved keys: {list(new_state.keys())}") + except Exception as e: + logger.warning(f"Could not preserve state keys: {e}") + + return rotated + + class TqdmLoggingHandler(logging.Handler): """Logging handler that writes through tqdm to avoid progress bar corruption.""" @@ -135,10 +185,10 @@ class NeoPig: self._items_since_save = 0 def _get_state_file(self, target_url: str) -> Path: - """Get path to state file for resume support.""" + """Get unified state file path for domain.""" parsed = urlparse(target_url) - domain = parsed.netloc.lower().replace('.', '-').replace(':', '-') - return self._state_dir / f"crawl-state-{domain}.json" + domain = parsed.netloc.lower() + return get_state_file_path(domain) def _save_state(self, target_url: str): """Save crawl state for resume.""" @@ -147,22 +197,29 @@ class NeoPig: return self._items_since_save = 0 - state = { + state_file = self._get_state_file(target_url) + + # Load existing state to preserve other keys (e.g., backfill state) + try: + full_state = json.loads(state_file.read_text()) if state_file.exists() else {} + except Exception: + full_state = {} + + # Update crawl section + full_state['crawl'] = { 'target_url': target_url, 'seen_media': list(self.seen_media), 'seen_screenshots': list(self.seen_screenshots), 'seen_pages': list(self.seen_pages), - # Note: skip_domains NOT persisted - domains may come back online - # It's saved during session for resume, but cleared on fresh runs 'skip_domains': list(self.fetcher.skip_domains), 'stats': self.stats, 'timestamp': datetime.now(timezone.utc).isoformat(), } + try: self._state_dir.mkdir(parents=True, exist_ok=True) - state_file = self._get_state_file(target_url) with open(state_file, 'w') as f: - json.dump(state, f) + json.dump(full_state, f, indent=2) except Exception as e: logger.debug(f"Failed to save state: {e}") @@ -178,18 +235,14 @@ class NeoPig: return False try: - with open(state_file, 'r') as f: - state = json.load(f) - # NOTE: Do NOT load seen_media or seen_screenshots from state file! - # Database is the source of truth - state file may have entries - # added before success (old buggy code). seen_media/seen_screenshots - # are loaded from DB in main() before crawl() is called. + full_state = json.loads(state_file.read_text()) + state = full_state.get('crawl', {}) + if not state: + return False # In fast mode, skip already-crawled pages for speed - # In normal mode, re-fetch pages to detect content changes (git handles versioning) if self.fast_mode: self.seen_pages = set(state.get('seen_pages', [])) - # Note: skip_domains NOT loaded - domains may have come back online saved_stats = state.get('stats', {}) for key in self.stats: if key in saved_stats: @@ -203,12 +256,15 @@ class NeoPig: logger.warning(f"Could not load state: {e}") return False - def _clear_state(self, target_url: str): - """Clear state file after successful completion.""" + def _clear_state(self, target_url: str, preserve_keys: List[str] = None): + """Rotate state file for fresh start. + + Args: + preserve_keys: Keys to preserve (e.g., ['crawl'] for backfill-only fresh) + """ try: state_file = self._get_state_file(target_url) - if state_file.exists(): - state_file.unlink() + rotate_state_file(state_file, preserve_keys=preserve_keys) except Exception: pass @@ -341,18 +397,19 @@ class NeoPig: self, url: str, md5_hash: str, + ext: str = 'jpg', ): """Create symlink in domain linkpeek vault pointing to hash vault.""" domain = self._get_domain(url) - # Domain screenshot path: vault/linkpeek_vault/{domain}/{url_path}.png + # Domain screenshot path: vault/linkpeek_vault/{domain}/{url_path}.{ext} parsed = urlparse(url) url_path = parsed.path.lstrip('/') or 'index' - url_path = url_path.replace('/', '_') + '.png' + url_path = url_path.replace('/', '_') + f'.{ext}' domain_ss_dir = Path(self.vault_path) / 'linkpeek_vault' / domain domain_ss_path = domain_ss_dir / url_path - # Hash vault path: vault/{hash[:2]}/{hash}.png - hash_vault_path = Path(self.vault_path) / md5_hash[:2] / f"{md5_hash}.png" + # Hash vault path: vault/{hash[:2]}/{hash}.{ext} + hash_vault_path = Path(self.vault_path) / md5_hash[:2] / f"{md5_hash}.{ext}" # Create symlink if not exists if not domain_ss_path.exists(): @@ -750,17 +807,19 @@ class NeoPig: md5_hash = result['md5_hash'] screenshot_data = result['data'] screenshot_size = len(screenshot_data) + screenshot_ext = result.get('format', 'png') + screenshot_mime = result.get('mime_type', 'image/png') # Screenshots are fetched by headless browser (network traffic) self.stats['bytes_downloaded'] += screenshot_size # Store in MD5 vault (for deduplication) if not await self.vault.exists(md5_hash): - await self.vault.store(md5_hash, screenshot_data, 'png') + await self.vault.store(md5_hash, screenshot_data, screenshot_ext) self.stats['bytes_stored'] += screenshot_size # Create symlink in linkpeek vault pointing to hash vault - await self._archive_screenshot_to_vault(page_uri, md5_hash) + await self._archive_screenshot_to_vault(page_uri, md5_hash, screenshot_ext) # Record in database as screenshot type await self.db.create_media_record( @@ -769,7 +828,7 @@ class NeoPig: page_uri=page_uri, crawl_job_id=job_id, media_type='screenshot', - mime_type='image/png', + mime_type=screenshot_mime, file_size=result.get('size', 0), page_title=page_title, page_description='', @@ -910,6 +969,222 @@ async def backfill_markdown(db_path: str, domain_filter: str = None, trim_wrappe logger.info(f"Backfill complete: {updated} pages updated") +async def backfill_screenshots( + db_path: str, + vault_path: str = "vault", + domain_filter: str = None, + delete_old: bool = True, + fast_mode: bool = False, +): + """Re-capture screenshots as JPEG to replace old PNGs. + + Args: + db_path: Path to SQLite database + vault_path: Path to vault directory + domain_filter: Only process pages matching this domain + delete_old: Delete old PNG files after successful JPEG capture + fast_mode: Skip delay between captures (for sites without robots.txt) + """ + from screenshot import ScreenshotCapture, ScreenshotConfig + from storage import ImageVault + from database import Database, Media, MediaSource, Page + from sqlalchemy import select, update + from datetime import datetime, timezone + import time + + # Load completed URIs from unified state file + completed_uris = set() + state_path = get_state_file_path(domain_filter) if domain_filter else None + if state_path and state_path.exists(): + try: + full_state = json.loads(state_path.read_text()) + completed_uris = set(full_state.get('backfill_screenshots', [])) + if completed_uris: + logger.info(f"Resuming: {len(completed_uris)} already backfilled") + except Exception as e: + logger.warning(f"Could not load state: {e}") + + config = ScreenshotConfig(enabled=True, full_page=True) + capture = ScreenshotCapture(config) + vault = ImageVault(vault_path) + await vault.init() + + db = Database(db_path) + await db.init() + + async with db.session() as session: + # Find all screenshot records using ORM + if domain_filter: + pattern = f"%://{domain_filter}%" if '://' not in domain_filter else f"%{domain_filter}%" + stmt = ( + select(MediaSource.page_uri, Media.md5_hash) + .join(Media, MediaSource.md5_hash == Media.md5_hash) + .where(Media.media_type == 'screenshot') + .where(MediaSource.page_uri.like(pattern)) + .distinct() + ) + logger.info(f"Backfilling screenshots for domain: {domain_filter}") + else: + stmt = ( + select(MediaSource.page_uri, Media.md5_hash) + .join(Media, MediaSource.md5_hash == Media.md5_hash) + .where(Media.media_type == 'screenshot') + .distinct() + ) + logger.info("Backfilling screenshots for ALL pages") + + result = await session.execute(stmt) + rows = result.fetchall() + total = len(rows) + logger.info(f"Found {total} screenshots to re-capture...") + + captured = 0 + skipped = 0 + failed = 0 + bytes_saved = 0 + domain_last_fetched = {} # Track last fetch time per domain + crawl_delay = 2.0 # Default crawl delay in seconds + + pbar = tqdm(rows, desc="Screenshots", unit="pages") + for row in pbar: + page_uri = row.page_uri + old_hash = row.md5_hash + domain = urlparse(page_uri).netloc.lower() + + # Skip if already completed (from state file) + if page_uri in completed_uris: + skipped += 1 + continue + + # Log current page + logger.info(f"Capturing: {page_uri}") + + try: + # Enforce crawl delay (skip in fast mode) + if not fast_mode: + last_fetched = domain_last_fetched.get(domain, 0) + elapsed = time.time() - last_fetched + if elapsed < crawl_delay: + await asyncio.sleep(crawl_delay - elapsed) + domain_last_fetched[domain] = time.time() + + # Get content length for dynamic delay calculation + from sqlalchemy import func + content_length_result = await session.execute( + select(func.length(Page.raw_html)).where(Page.uri == page_uri) + ) + content_length = content_length_result.scalar() or 0 + + # Capture new screenshot (will be JPEG) with dynamic delay for large pages + result = await capture.capture(page_uri, content_length=content_length) + if not result: + failed += 1 + continue + + new_hash = result['md5_hash'] + new_data = result['data'] + new_ext = result.get('format', 'jpg') + new_mime = result.get('mime_type', 'image/jpeg') + new_size = len(new_data) + + # Get old file size for comparison + old_path = Path(vault_path) / old_hash[:2] / f"{old_hash}.png" + old_size = old_path.stat().st_size if old_path.exists() else 0 + + # Store new screenshot + if not await vault.exists(new_hash): + await vault.store(new_hash, new_data, new_ext) + + # Delete old PNG file FIRST (before metadata update) + if delete_old and old_hash != new_hash and old_path.exists(): + old_path.unlink() + bytes_saved += old_size - new_size + + # Update linkpeek symlink + parsed = urlparse(page_uri) + url_path = parsed.path.lstrip('/') or 'index' + + # Remove old symlink + old_symlink = Path(vault_path) / 'linkpeek_vault' / domain / (url_path.replace('/', '_') + '.png') + if old_symlink.exists(): + old_symlink.unlink() + + # Create new symlink + new_symlink = Path(vault_path) / 'linkpeek_vault' / domain / (url_path.replace('/', '_') + f'.{new_ext}') + new_symlink.parent.mkdir(parents=True, exist_ok=True) + hash_vault_path = Path(vault_path) / new_hash[:2] / f"{new_hash}.{new_ext}" + rel_path = os.path.relpath(hash_vault_path, new_symlink.parent) + if not new_symlink.exists(): + new_symlink.symlink_to(rel_path) + + # Update database using ORM + # md5_hash is PRIMARY KEY, so we need to: insert new -> update refs -> delete old + from sqlalchemy import delete + + # Get old record data + old_media_result = await session.execute( + select(Media).where(Media.md5_hash == old_hash) + ) + old_media = old_media_result.scalar_one_or_none() + + if old_media and old_hash != new_hash: + # Insert new media record + new_media = Media( + md5_hash=new_hash, + media_type='screenshot', + mime_type=new_mime, + file_size=new_size, + keywords=old_media.keywords, + alt_text=old_media.alt_text, + title=old_media.title, + first_seen_at=datetime.now(timezone.utc).isoformat(), + analysis_status=old_media.analysis_status, + analysis_result=old_media.analysis_result, + ) + session.add(new_media) + await session.flush() + + # Update MediaSource references + await session.execute( + update(MediaSource) + .where(MediaSource.md5_hash == old_hash) + .values(md5_hash=new_hash) + ) + + # Delete old media record + await session.execute( + delete(Media).where(Media.md5_hash == old_hash) + ) + elif old_media: + # Same hash, just update metadata + old_media.mime_type = new_mime + old_media.file_size = new_size + old_media.first_seen_at = datetime.now(timezone.utc).isoformat() + + captured += 1 + # Commit after each capture to preserve progress + await session.commit() + + # Save state after each successful capture + if state_path: + completed_uris.add(page_uri) + try: + full_state = json.loads(state_path.read_text()) if state_path.exists() else {} + except Exception: + full_state = {} + full_state['backfill_screenshots'] = list(completed_uris) + state_path.write_text(json.dumps(full_state, indent=2)) + + except Exception as e: + logger.warning(f"Error re-capturing {page_uri}: {e}") + failed += 1 + + await session.commit() + logger.info(f"Backfill complete: {captured} captured, {skipped} skipped, {failed} failed") + if bytes_saved > 0: + logger.info(f"Space saved: {bytes_saved / 1024 / 1024:.1f} MB") + + async def main(): parser = argparse.ArgumentParser( description="neopig - Neo Python Image Grabber", @@ -1018,7 +1293,7 @@ async def main(): parser.add_argument( "--fresh", action="store_true", - help="Start fresh, ignoring any saved resume state" + help="Start fresh (rotates state file; backfills preserve crawl state)" ) parser.add_argument( @@ -1039,6 +1314,18 @@ async def main(): help="With --backfill-markdown: strip nav/header/footer/logo before conversion" ) + parser.add_argument( + "--backfill-screenshots", + metavar="DOMAIN", + help="Re-capture screenshots as JPEG for DOMAIN (e.g., example.com)" + ) + + parser.add_argument( + "--keep-old-screenshots", + action="store_true", + help="With --backfill-screenshots: keep old PNG files instead of deleting them" + ) + parser.add_argument( "--serve", action="store_true", @@ -1079,6 +1366,39 @@ async def main(): await backfill_markdown(args.db, domain_filter=args.backfill_markdown, trim_wrapper=args.trim_wrapper) return + # Handle --backfill-screenshots + if args.backfill_screenshots: + # Handle --fresh: rotate state but preserve crawl data + if args.fresh: + state_path = get_state_file_path(args.backfill_screenshots) + rotate_state_file(state_path, preserve_keys=['crawl']) + + serp_process = None + if args.serve: + import subprocess, sys + serp_script = Path(__file__).parent / 'serp.py' + if serp_script.exists(): + serp_cmd = [ + sys.executable, str(serp_script), + '--port', str(args.port), + '--db', args.db, + '--vault', args.vault, + ] + serp_process = subprocess.Popen(serp_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + logger.info(f"SERP server started at http://localhost:{args.port}") + try: + await backfill_screenshots( + args.db, + vault_path=args.vault, + domain_filter=args.backfill_screenshots, + delete_old=not args.keep_old_screenshots, + fast_mode=args.fast, + ) + finally: + if serp_process: + serp_process.terminate() + return + # Require targets for crawling if not args.targets: parser.error("targets required (use --list-engines to see available screenshot engines)") diff --git a/screenshot.py b/screenshot.py index 1123459..5b4e980 100644 --- a/screenshot.py +++ b/screenshot.py @@ -19,6 +19,7 @@ or you can specify an engine explicitly. import asyncio import hashlib +import io import logging import shutil from dataclasses import dataclass, field @@ -26,6 +27,8 @@ from pathlib import Path from typing import Optional, List, Dict, Any import tempfile +from PIL import Image + logger = logging.getLogger(__name__) # Engine preference order - lightest/fastest first @@ -49,13 +52,15 @@ NATIVE_ENGINES = {'wkhtmltoimage', 'cutycapt'} class ScreenshotConfig: """Screenshot capture configuration.""" enabled: bool = False - width: int = 1280 - height: int = 1024 + width: int = 1024 + height: int = 768 delay: int = 1000 # ms after DOM load timeout: int = 30000 # ms total timeout user_agent: Optional[str] = None engine: Optional[str] = None # None = auto-detect best available full_page: bool = False + format: str = 'jpeg' # 'jpeg' or 'png' + quality: int = 93 # JPEG quality (1-100) class ScreenshotCapture: @@ -185,13 +190,13 @@ class ScreenshotCapture: content_length: Length of raw HTML content in bytes Returns: - Delay in milliseconds (min 1000ms, max 15000ms) + Delay in milliseconds (min 2000ms, max 30000ms) """ - # Base delay of 1 second - base_delay = 1000 + # Base delay of 2 seconds + base_delay = 2000 - # Add 1 second per 50KB of content, up to a max - additional_delay = min((content_length // 50000) * 1000, 14000) + # Add 2 seconds per 50KB of content, up to a max of 28 seconds + additional_delay = min((content_length // 50000) * 2000, 28000) return base_delay + additional_delay @@ -257,21 +262,38 @@ class ScreenshotCapture: return None # Read bytes from output file - data = Path(output_path).read_bytes() - if not data: + png_data = Path(output_path).read_bytes() + if not png_data: logger.warning(f"Screenshot empty: {uri}") return None + # Convert to JPEG if configured + if self.config.format == 'jpeg': + img = Image.open(io.BytesIO(png_data)) + # Convert RGBA to RGB (JPEG doesn't support alpha) + if img.mode in ('RGBA', 'LA', 'P'): + img = img.convert('RGB') + output = io.BytesIO() + img.save(output, format='JPEG', quality=self.config.quality, optimize=True) + data = output.getvalue() + mime_type = 'image/jpeg' + ext = 'jpg' + else: + data = png_data + mime_type = 'image/png' + ext = 'png' + md5_hash = hashlib.md5(data).hexdigest() - logger.debug(f"Screenshot captured ({self._engine_name}): {uri} -> {md5_hash}") + logger.debug(f"Screenshot captured ({self._engine_name}, {ext}): {uri} -> {md5_hash} ({len(data)//1024}KB)") return { 'data': data, 'md5_hash': md5_hash, - 'mime_type': 'image/png', + 'mime_type': mime_type, 'size': len(data), 'source_uri': uri, 'engine': self._engine_name, + 'format': ext, } finally: # Clean up temp file