dna-sequencing-match

"""
DNA Sequence Search System
Searches for gene sequences in complementary DNA strands stored in dna sequence 6GB
simulate DB with connection pooling (max 10 connections per database)

"""

from __future__ import annotations

import argparse
import os
import time
from typing import Iterable, List, Optional


def _human_readable(n: int) -> str:
	for unit in ('B', 'KB', 'MB', 'GB', 'TB'):
		if n < 1024:
			return f"{n:.2f}{unit}"
		n /= 1024
	return f"{n:.2f}PB"


def reverse_complement(seq: bytes) -> bytes:
	"""Return the reverse-complement of a DNA sequence bytes (A,T,G,C)."""
	table = bytes.maketrans(b'ATGCatgc', b'TACGtacg')
	return seq.translate(table)[::-1]


def normalize_pattern(p: str) -> bytes:
	return p.strip().upper().encode('ascii')


def stream_search(file_path: str, pattern: bytes, chunk_size: int = 1 << 20) -> Iterable[int]:
	"""Yield file offsets where `pattern` occurs in `file_path` using streaming reads.

	This is memory efficient and handles matches that span chunk boundaries by
	carrying an overlap of len(pattern)-1 bytes between chunks.
	"""
	if not pattern:
		return

	patlen = len(pattern)
	overlap = patlen - 1

	with open(file_path, 'rb') as f:
		pos = 0
		tail = b''
		while True:
			chunk = f.read(chunk_size)
			if not chunk:
				break
			buf = tail + chunk
			start = 0
			while True:
				idx = buf.find(pattern, start)
				if idx == -1:
					break
				yield pos - len(tail) + idx
				start = idx + 1

			# prepare tail for next chunk
			if overlap > 0:
				tail = buf[-overlap:]
			else:
				tail = b''
			pos += len(chunk)


def main(argv: Optional[List[str]] = None) -> int:
	parser = argparse.ArgumentParser(description='Streaming DNA sequence search')
	parser.add_argument('--file', '-f', help='Path to DNA sequence file', required=False)
	parser.add_argument('--pattern', '-p', help='DNA pattern to search for (e.g. ATG)', required=False)
	parser.add_argument('--pattern-file', help='File containing pattern (first line used)')
	parser.add_argument('--reverse', '-r', action='store_true', help='Also search reverse-complement')
	parser.add_argument('--chunk-size', type=int, default=1 << 20, help='Read chunk size in bytes')
	parser.add_argument('--max-results', type=int, default=0, help='Stop after this many results (0 = unlimited)')
	parser.add_argument('--example', action='store_true', help='Create a tiny example file and run a sample search')

	args = parser.parse_args(argv)

	if args.example:
		sample = b'AAATGCCCATGTTTATGCGTATG'
		sample_path = os.path.join(os.getcwd(), 'sample_dna.txt')
		with open(sample_path, 'wb') as f:
			f.write(sample)
		size = os.path.getsize(sample_path)
		print(f'Wrote example file: {sample_path} (size: {_human_readable(size)})')
		args.file = sample_path
		if not args.pattern:
			args.pattern = 'ATG'

	if not args.file:
		parser.error('--file is required (or use --example)')

	if args.pattern_file:
		with open(args.pattern_file, 'r', encoding='utf-8') as pf:
			args.pattern = pf.readline().strip()

	if not args.pattern:
		parser.error('--pattern is required (or use --pattern-file)')

	pattern = normalize_pattern(args.pattern)
	if any(c not in b'ATGC' for c in pattern):
		print('Warning: pattern contains non-ATGC characters; results may be unexpected')

	matches_found = 0
	file_size = os.path.getsize(args.file) if os.path.exists(args.file) else 0
	print(f"Searching {args.file} for pattern {pattern.decode()} (chunk={args.chunk_size}, size={_human_readable(file_size)})")
	start_time = time.time()

	for off in stream_search(args.file, pattern, chunk_size=args.chunk_size):
		matches_found += 1
		if args.max_results and matches_found >= args.max_results:
			break

	if args.reverse:
		rc = reverse_complement(pattern)
		print(f"Searching reverse-complement: {rc.decode()}")
		for off in stream_search(args.file, rc, chunk_size=args.chunk_size):
			matches_found += 1
			if args.max_results and matches_found >= args.max_results:
				break

	duration = time.time() - start_time
	if duration > 0:
		rate = matches_found / duration
		print(f"Search complete: {matches_found} matches in {duration:.2f}s ({rate:.2f} matches/s)")
	else:
		print(f"Search complete: {matches_found} matches (duration too short to measure)")
	return 0


if __name__ == '__main__':
	raise SystemExit(main())