engineering/pig.py
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

modified: .gitignore

Browse source 
modified:   archive.py
	modified:   html2md.py
	modified:   neopig.py
	modified:   screenshot.py
This commit is contained in:
Russell Ballestrini 2025-12-30 09:50:54 -05:00
parent 623213f5b6
commit 7b4c8b0241
5 changed files with 529 additions and 115 deletions
Download patch file Download diff file Expand all files Collapse all files

215
archive.py
View file

@ -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