import os
import tempfile
import random
from typing import Optional
class DNA_SEQUENCE_FILE:
"""Simple file-backed DNA sequence helper.
Stores ASCII bytes 'A','T','G','C'. Provides chunked random-generation,
read, write (overwrite/append/at-offset), and insertion (streaming via temp file).
"""
BASES = b'ATGC'
def __init__(self, filename: str, filesize: int = 0):
self.filename = filename
self.filesize = int(filesize)
def __str__(self) -> str:
return f"DNA_SEQUENCE_FILE(filename={self.filename!r}, filesize={self.filesize})"
def _fsync_file(self, f) -> None:
f.flush()
try:
os.fsync(f.fileno())
except OSError:
# fallback/no-op on platforms that don't support fsync
pass
def read(self, offset: int = 0, length: Optional[int] = None) -> bytes:
"""Read up to `length` bytes from file starting at `offset`.
If `length` is None, read to EOF.
"""
with open(self.filename, 'rb') as f:
f.seek(offset)
data = f.read() if length is None else f.read(length)
return data
def write(self, data: bytes, offset: Optional[int] = None) -> None:
"""Write `data` to file. If `offset` is None, overwrite file; else write at offset.
Use binary-safe operations and fsync after writing.
"""
if offset is None:
# overwrite entire file
with open(self.filename, 'wb') as f:
f.write(data)
self._fsync_file(f)
self.filesize = os.path.getsize(self.filename)
else:
# write in-place at offset (requires file exists or will be created)
os.makedirs(os.path.dirname(self.filename) or '.', exist_ok=True)
# open r+b if file exists else create then r+b
mode = 'r+b' if os.path.exists(self.filename) else 'w+b'
with open(self.filename, mode) as f:
f.seek(offset)
f.write(data)
self._fsync_file(f)
self.filesize = os.path.getsize(self.filename)
def insert(self, offset: int, data: bytes) -> None:
"""Insert `data` at byte `offset` by streaming into a temporary file.
For large files this is IO heavy but safe and memory-efficient.
"""
# If file doesn't exist, insertion is just a write
if not os.path.exists(self.filename):
self.write(data, offset=0)
return
with open(self.filename, 'rb') as src:
# Create temp file in same directory to allow atomic replace
dirpath = os.path.dirname(self.filename) or '.'
fd, tmp_path = tempfile.mkstemp(dir=dirpath)
with os.fdopen(fd, 'wb') as dst:
# copy up to offset
to_copy = offset
chunk = 1 << 20
while to_copy > 0:
read_len = min(chunk, to_copy)
chunk_data = src.read(read_len)
if not chunk_data:
break
dst.write(chunk_data)
to_copy -= len(chunk_data)
# write inserted data
dst.write(data)
# copy remainder
while True:
chunk_data = src.read(chunk)
if not chunk_data:
break
dst.write(chunk_data)
self._fsync_file(dst)
# replace original file atomically
os.replace(tmp_path, self.filename)
self.filesize = os.path.getsize(self.filename)
def generate_random(self, size: int, chunk_size: int = 1 << 20, seed: Optional[int] = None) -> None:
"""Generate a random DNA sequence of exactly `size` bytes and write to the file.
Generation is chunked to avoid large memory usage. Each output byte is one of 'A','T','G','C'.
"""
if seed is not None:
random.seed(seed)
bases = self.BASES
# Precompute a mapping for 0-255 -> base using modulo for speed
table = bytes(bytes([bases[b % 4] for b in range(256)]))
with open(self.filename, 'wb') as f:
remaining = int(size)
while remaining > 0:
this = min(remaining, chunk_size)
# get cryptographic random bytes and map via table
rnd = os.urandom(this)
out = rnd.translate(table)
f.write(out)
remaining -= this
self._fsync_file(f)
self.filesize = os.path.getsize(self.filename)
def insert_random_bases(self, count: int, seed: Optional[int] = None) -> None:
"""Insert `count` random single-base bytes at random offsets throughout the file.
This performs `count` insertions; for large counts this may be slow.
"""
if seed is not None:
random.seed(seed)
for _ in range(count):
base = random.choice(self.BASES)
# choose offset in [0, filesize]
off = random.randint(0, max(0, self.filesize))
self.insert(off, bytes([base]))
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"
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='DNA sequence file utility (safe, chunked)')
parser.add_argument('file', help='path to sequence file to create or operate on')
parser.add_argument('--size', type=str, default=None,
help='generate random sequence of given size, e.g. 6GB, 1048576B')
parser.add_argument('--generate', action='store_true', help='generate random sequence of given size')
parser.add_argument('--insert-random', type=int, default=0, help='insert N random bases at random offsets')
parser.add_argument('--seed', type=int, default=None, help='seed for deterministic generation')
args = parser.parse_args()
def parse_size(size: str) -> int:
"""Given a human-readable byte string (e.g. 2G, 10GB, 30MB, 20KB),
return the number of bytes. Will return 0 if the argument has
unexpected form.
from https://gist.github.com/beugley/ccd69945346759eb6142272a6d69b4e0
"""
if (size[-1] == 'B'):
size = size[:-1]
if (size.isdigit()):
bytes = int(size)
else:
bytes = size[:-1]
unit = size[-1]
if (bytes.isdigit()):
bytes = int(bytes)
if (unit == 'G'):
bytes *= 1073741824
elif (unit == 'M'):
bytes *= 1048576
elif (unit == 'K'):
bytes *= 1024
else:
bytes = 0
else:
bytes = 0
return bytes
obj = DNA_SEQUENCE_FILE(args.file)
if args.generate:
if not args.size:
parser.error('--generate requires --size')
size_bytes = parse_size(args.size)
print(f'Generating {size_bytes} bytes ({_human_readable(size_bytes)}) into {args.file}')
obj.generate_random(size_bytes, seed=args.seed)
print('Done:', obj)
if args.insert_random:
print(f'Inserting {args.insert_random} random bases into {args.file}')
obj.filesize = os.path.getsize(args.file) if os.path.exists(args.file) else 0
obj.insert_random_bases(args.insert_random, seed=args.seed)
print('Done:', obj)