Mercurial > repos > rliterman > csp2

annotate CSP2/CSP2_env/env-d9b9114564458d9d-741b3de822f2aaca6c6caa4325c4afce/lib/python3.8/site-packages/Bio/File.py @ 68:5028fdace37b

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload commit 2e9511a184a1ca667c7be0c6321a36dc4e3d116d
author jpayne
date Tue, 18 Mar 2025 16:23:26 -0400
parents
children
rev   line source
jpayne@68 1 # Copyright 1999 by Jeffrey Chang. All rights reserved.
jpayne@68 2 # Copyright 2009-2018 by Peter Cock. All rights reserved.
jpayne@68 3 #
jpayne@68 4 # This file is part of the Biopython distribution and governed by your
jpayne@68 5 # choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
jpayne@68 6 # Please see the LICENSE file that should have been included as part of this
jpayne@68 7 # package.
jpayne@68 8 """Code for more fancy file handles.
jpayne@68 9
jpayne@68 10 Bio.File defines private classes used in Bio.SeqIO and Bio.SearchIO for
jpayne@68 11 indexing files. These are not intended for direct use.
jpayne@68 12 """
jpayne@68 13
jpayne@68 14 import os
jpayne@68 15 import contextlib
jpayne@68 16 import itertools
jpayne@68 17 import collections.abc
jpayne@68 18
jpayne@68 19 from abc import ABC, abstractmethod
jpayne@68 20
jpayne@68 21 try:
jpayne@68 22 import sqlite3
jpayne@68 23 except ImportError:
jpayne@68 24 # May be missing if Python was compiled from source without its dependencies
jpayne@68 25 sqlite3 = None # type: ignore
jpayne@68 26
jpayne@68 27
jpayne@68 28 @contextlib.contextmanager
jpayne@68 29 def as_handle(handleish, mode="r", **kwargs):
jpayne@68 30 r"""Context manager to ensure we are using a handle.
jpayne@68 31
jpayne@68 32 Context manager for arguments that can be passed to SeqIO and AlignIO read, write,
jpayne@68 33 and parse methods: either file objects or path-like objects (strings, pathlib.Path
jpayne@68 34 instances, or more generally, anything that can be handled by the builtin 'open'
jpayne@68 35 function).
jpayne@68 36
jpayne@68 37 When given a path-like object, returns an open file handle to that path, with provided
jpayne@68 38 mode, which will be closed when the manager exits.
jpayne@68 39
jpayne@68 40 All other inputs are returned, and are *not* closed.
jpayne@68 41
jpayne@68 42 Arguments:
jpayne@68 43 - handleish - Either a file handle or path-like object (anything which can be
jpayne@68 44 passed to the builtin 'open' function, such as str, bytes,
jpayne@68 45 pathlib.Path, and os.DirEntry objects)
jpayne@68 46 - mode - Mode to open handleish (used only if handleish is a string)
jpayne@68 47 - kwargs - Further arguments to pass to open(...)
jpayne@68 48
jpayne@68 49 Examples
jpayne@68 50 --------
jpayne@68 51 >>> from Bio import File
jpayne@68 52 >>> import os
jpayne@68 53 >>> with File.as_handle('seqs.fasta', 'w') as fp:
jpayne@68 54 ... fp.write('>test\nACGT')
jpayne@68 55 ...
jpayne@68 56 10
jpayne@68 57 >>> fp.closed
jpayne@68 58 True
jpayne@68 59
jpayne@68 60 >>> handle = open('seqs.fasta', 'w')
jpayne@68 61 >>> with File.as_handle(handle) as fp:
jpayne@68 62 ... fp.write('>test\nACGT')
jpayne@68 63 ...
jpayne@68 64 10
jpayne@68 65 >>> fp.closed
jpayne@68 66 False
jpayne@68 67 >>> fp.close()
jpayne@68 68 >>> os.remove("seqs.fasta") # tidy up
jpayne@68 69
jpayne@68 70 """
jpayne@68 71 try:
jpayne@68 72 with open(handleish, mode, **kwargs) as fp:
jpayne@68 73 yield fp
jpayne@68 74 except TypeError:
jpayne@68 75 yield handleish
jpayne@68 76
jpayne@68 77
jpayne@68 78 def _open_for_random_access(filename):
jpayne@68 79 """Open a file in binary mode, spot if it is BGZF format etc (PRIVATE).
jpayne@68 80
jpayne@68 81 This functionality is used by the Bio.SeqIO and Bio.SearchIO index
jpayne@68 82 and index_db functions.
jpayne@68 83
jpayne@68 84 If the file is gzipped but not BGZF, a specific ValueError is raised.
jpayne@68 85 """
jpayne@68 86 handle = open(filename, "rb")
jpayne@68 87 magic = handle.read(2)
jpayne@68 88 handle.seek(0)
jpayne@68 89
jpayne@68 90 if magic == b"\x1f\x8b":
jpayne@68 91 # This is a gzipped file, but is it BGZF?
jpayne@68 92 from . import bgzf
jpayne@68 93
jpayne@68 94 try:
jpayne@68 95 # If it is BGZF, we support that
jpayne@68 96 return bgzf.BgzfReader(mode="rb", fileobj=handle)
jpayne@68 97 except ValueError as e:
jpayne@68 98 assert "BGZF" in str(e)
jpayne@68 99 # Not a BGZF file after all,
jpayne@68 100 handle.close()
jpayne@68 101 raise ValueError(
jpayne@68 102 "Gzipped files are not suitable for indexing, "
jpayne@68 103 "please use BGZF (blocked gzip format) instead."
jpayne@68 104 ) from None
jpayne@68 105
jpayne@68 106 return handle
jpayne@68 107
jpayne@68 108
jpayne@68 109 # The rest of this file defines code used in Bio.SeqIO and Bio.SearchIO
jpayne@68 110 # for indexing
jpayne@68 111
jpayne@68 112
jpayne@68 113 class _IndexedSeqFileProxy(ABC):
jpayne@68 114 """Abstract base class for file format specific random access (PRIVATE).
jpayne@68 115
jpayne@68 116 This is subclasses in both Bio.SeqIO for indexing as SeqRecord
jpayne@68 117 objects, and in Bio.SearchIO for indexing QueryResult objects.
jpayne@68 118
jpayne@68 119 Subclasses for each file format should define '__iter__', 'get'
jpayne@68 120 and optionally 'get_raw' methods.
jpayne@68 121 """
jpayne@68 122
jpayne@68 123 @abstractmethod
jpayne@68 124 def __iter__(self):
jpayne@68 125 """Return (identifier, offset, length in bytes) tuples.
jpayne@68 126
jpayne@68 127 The length can be zero where it is not implemented or not
jpayne@68 128 possible for a particular file format.
jpayne@68 129 """
jpayne@68 130 raise NotImplementedError
jpayne@68 131
jpayne@68 132 @abstractmethod
jpayne@68 133 def get(self, offset):
jpayne@68 134 """Return parsed object for this entry."""
jpayne@68 135 # Most file formats with self contained records can be handled by
jpayne@68 136 # parsing StringIO(self.get_raw(offset).decode())
jpayne@68 137 raise NotImplementedError
jpayne@68 138
jpayne@68 139 def get_raw(self, offset):
jpayne@68 140 """Return the raw record from the file as a bytes string (if implemented).
jpayne@68 141
jpayne@68 142 If the key is not found, a KeyError exception is raised.
jpayne@68 143
jpayne@68 144 This may not have been implemented for all file formats.
jpayne@68 145 """
jpayne@68 146 # Should be done by each sub-class (if possible)
jpayne@68 147 raise NotImplementedError("Not available for this file format.")
jpayne@68 148
jpayne@68 149
jpayne@68 150 class _IndexedSeqFileDict(collections.abc.Mapping):
jpayne@68 151 """Read only dictionary interface to a sequential record file.
jpayne@68 152
jpayne@68 153 This code is used in both Bio.SeqIO for indexing as SeqRecord
jpayne@68 154 objects, and in Bio.SearchIO for indexing QueryResult objects.
jpayne@68 155
jpayne@68 156 Keeps the keys and associated file offsets in memory, reads the file
jpayne@68 157 to access entries as objects parsing them on demand. This approach
jpayne@68 158 is memory limited, but will work even with millions of records.
jpayne@68 159
jpayne@68 160 Note duplicate keys are not allowed. If this happens, a ValueError
jpayne@68 161 exception is raised.
jpayne@68 162
jpayne@68 163 As used in Bio.SeqIO, by default the SeqRecord's id string is used
jpayne@68 164 as the dictionary key. In Bio.SearchIO, the query's id string is
jpayne@68 165 used. This can be changed by supplying an optional key_function,
jpayne@68 166 a callback function which will be given the record id and must
jpayne@68 167 return the desired key. For example, this allows you to parse
jpayne@68 168 NCBI style FASTA identifiers, and extract the GI number to use
jpayne@68 169 as the dictionary key.
jpayne@68 170
jpayne@68 171 Note that this dictionary is essentially read only. You cannot
jpayne@68 172 add or change values, pop values, nor clear the dictionary.
jpayne@68 173 """
jpayne@68 174
jpayne@68 175 def __init__(self, random_access_proxy, key_function, repr, obj_repr):
jpayne@68 176 """Initialize the class."""
jpayne@68 177 # Use key_function=None for default value
jpayne@68 178 self._proxy = random_access_proxy
jpayne@68 179 self._key_function = key_function
jpayne@68 180 self._repr = repr
jpayne@68 181 self._obj_repr = obj_repr
jpayne@68 182 self._cached_prev_record = (None, None) # (key, record)
jpayne@68 183 if key_function:
jpayne@68 184 offset_iter = (
jpayne@68 185 (key_function(key), offset, length)
jpayne@68 186 for (key, offset, length) in random_access_proxy
jpayne@68 187 )
jpayne@68 188 else:
jpayne@68 189 offset_iter = random_access_proxy
jpayne@68 190 offsets = {}
jpayne@68 191 for key, offset, length in offset_iter:
jpayne@68 192 # Note - we don't store the length because I want to minimise the
jpayne@68 193 # memory requirements. With the SQLite backend the length is kept
jpayne@68 194 # and is used to speed up the get_raw method (by about 3 times).
jpayne@68 195 # The length should be provided by all the current backends except
jpayne@68 196 # SFF where there is an existing Roche index we can reuse (very fast
jpayne@68 197 # but lacks the record lengths)
jpayne@68 198 # assert length or format in ["sff", "sff-trim"], \
jpayne@68 199 # "%s at offset %i given length %r (%s format %s)" \
jpayne@68 200 # % (key, offset, length, filename, format)
jpayne@68 201 if key in offsets:
jpayne@68 202 self._proxy._handle.close()
jpayne@68 203 raise ValueError(f"Duplicate key '{key}'")
jpayne@68 204 else:
jpayne@68 205 offsets[key] = offset
jpayne@68 206 self._offsets = offsets
jpayne@68 207
jpayne@68 208 def __repr__(self):
jpayne@68 209 """Return a string representation of the File object."""
jpayne@68 210 return self._repr
jpayne@68 211
jpayne@68 212 def __str__(self):
jpayne@68 213 """Create a string representation of the File object."""
jpayne@68 214 # TODO - How best to handle the __str__ for SeqIO and SearchIO?
jpayne@68 215 if self:
jpayne@68 216 return f"{{{list(self.keys())[0]!r} : {self._obj_repr}(...), ...}}"
jpayne@68 217 else:
jpayne@68 218 return "{}"
jpayne@68 219
jpayne@68 220 def __len__(self):
jpayne@68 221 """Return the number of records."""
jpayne@68 222 return len(self._offsets)
jpayne@68 223
jpayne@68 224 def __iter__(self):
jpayne@68 225 """Iterate over the keys."""
jpayne@68 226 return iter(self._offsets)
jpayne@68 227
jpayne@68 228 def __getitem__(self, key):
jpayne@68 229 """Return record for the specified key.
jpayne@68 230
jpayne@68 231 As an optimization when repeatedly asked to look up the same record,
jpayne@68 232 the key and record are cached so that if the *same* record is
jpayne@68 233 requested next time, it can be returned without going to disk.
jpayne@68 234 """
jpayne@68 235 if key == self._cached_prev_record[0]:
jpayne@68 236 return self._cached_prev_record[1]
jpayne@68 237 # Pass the offset to the proxy
jpayne@68 238 record = self._proxy.get(self._offsets[key])
jpayne@68 239 if self._key_function:
jpayne@68 240 key2 = self._key_function(record.id)
jpayne@68 241 else:
jpayne@68 242 key2 = record.id
jpayne@68 243 if key != key2:
jpayne@68 244 raise ValueError(f"Key did not match ({key} vs {key2})")
jpayne@68 245 self._cached_prev_record = (key, record)
jpayne@68 246 return record
jpayne@68 247
jpayne@68 248 def get_raw(self, key):
jpayne@68 249 """Return the raw record from the file as a bytes string.
jpayne@68 250
jpayne@68 251 If the key is not found, a KeyError exception is raised.
jpayne@68 252 """
jpayne@68 253 # Pass the offset to the proxy
jpayne@68 254 return self._proxy.get_raw(self._offsets[key])
jpayne@68 255
jpayne@68 256 def close(self):
jpayne@68 257 """Close the file handle being used to read the data.
jpayne@68 258
jpayne@68 259 Once called, further use of the index won't work. The sole purpose
jpayne@68 260 of this method is to allow explicit handle closure - for example
jpayne@68 261 if you wish to delete the file, on Windows you must first close
jpayne@68 262 all open handles to that file.
jpayne@68 263 """
jpayne@68 264 self._proxy._handle.close()
jpayne@68 265
jpayne@68 266
jpayne@68 267 class _SQLiteManySeqFilesDict(_IndexedSeqFileDict):
jpayne@68 268 """Read only dictionary interface to many sequential record files.
jpayne@68 269
jpayne@68 270 This code is used in both Bio.SeqIO for indexing as SeqRecord
jpayne@68 271 objects, and in Bio.SearchIO for indexing QueryResult objects.
jpayne@68 272
jpayne@68 273 Keeps the keys, file-numbers and offsets in an SQLite database. To access
jpayne@68 274 a record by key, reads from the offset in the appropriate file and then
jpayne@68 275 parses the record into an object.
jpayne@68 276
jpayne@68 277 There are OS limits on the number of files that can be open at once,
jpayne@68 278 so a pool are kept. If a record is required from a closed file, then
jpayne@68 279 one of the open handles is closed first.
jpayne@68 280 """
jpayne@68 281
jpayne@68 282 def __init__(
jpayne@68 283 self,
jpayne@68 284 index_filename,
jpayne@68 285 filenames,
jpayne@68 286 proxy_factory,
jpayne@68 287 fmt,
jpayne@68 288 key_function,
jpayne@68 289 repr,
jpayne@68 290 max_open=10,
jpayne@68 291 ):
jpayne@68 292 """Initialize the class."""
jpayne@68 293 # TODO? - Don't keep filename list in memory (just in DB)?
jpayne@68 294 # Should save a chunk of memory if dealing with 1000s of files.
jpayne@68 295 # Furthermore could compare a generator to the DB on reloading
jpayne@68 296 # (no need to turn it into a list)
jpayne@68 297
jpayne@68 298 if sqlite3 is None:
jpayne@68 299 # Python was compiled without sqlite3 support
jpayne@68 300 from Bio import MissingPythonDependencyError
jpayne@68 301
jpayne@68 302 raise MissingPythonDependencyError(
jpayne@68 303 "Python was compiled without the sqlite3 module"
jpayne@68 304 )
jpayne@68 305 if filenames is not None:
jpayne@68 306 filenames = list(filenames) # In case it was a generator
jpayne@68 307
jpayne@68 308 # Cache the arguments as private variables
jpayne@68 309 self._index_filename = index_filename
jpayne@68 310 self._filenames = filenames
jpayne@68 311 self._format = fmt
jpayne@68 312 self._key_function = key_function
jpayne@68 313 self._proxy_factory = proxy_factory
jpayne@68 314 self._repr = repr
jpayne@68 315 self._max_open = max_open
jpayne@68 316 self._proxies = {}
jpayne@68 317
jpayne@68 318 # Note if using SQLite :memory: trick index filename, this will
jpayne@68 319 # give $PWD as the relative path (which is fine).
jpayne@68 320 self._relative_path = os.path.abspath(os.path.dirname(index_filename))
jpayne@68 321
jpayne@68 322 if os.path.isfile(index_filename):
jpayne@68 323 self._load_index()
jpayne@68 324 else:
jpayne@68 325 self._build_index()
jpayne@68 326
jpayne@68 327 def _load_index(self):
jpayne@68 328 """Call from __init__ to re-use an existing index (PRIVATE)."""
jpayne@68 329 index_filename = self._index_filename
jpayne@68 330 relative_path = self._relative_path
jpayne@68 331 filenames = self._filenames
jpayne@68 332 fmt = self._format
jpayne@68 333 proxy_factory = self._proxy_factory
jpayne@68 334
jpayne@68 335 con = sqlite3.dbapi2.connect(index_filename, check_same_thread=False)
jpayne@68 336 self._con = con
jpayne@68 337 # Check the count...
jpayne@68 338 try:
jpayne@68 339 (count,) = con.execute(
jpayne@68 340 "SELECT value FROM meta_data WHERE key=?;", ("count",)
jpayne@68 341 ).fetchone()
jpayne@68 342 self._length = int(count)
jpayne@68 343 if self._length == -1:
jpayne@68 344 con.close()
jpayne@68 345 raise ValueError("Unfinished/partial database") from None
jpayne@68 346
jpayne@68 347 # use MAX(_ROWID_) to obtain the number of sequences in the database
jpayne@68 348 # using COUNT(key) is quite slow in SQLITE
jpayne@68 349 # (https://stackoverflow.com/questions/8988915/sqlite-count-slow-on-big-tables)
jpayne@68 350 (count,) = con.execute("SELECT MAX(_ROWID_) FROM offset_data;").fetchone()
jpayne@68 351 if self._length != int(count):
jpayne@68 352 con.close()
jpayne@68 353 raise ValueError(
jpayne@68 354 "Corrupt database? %i entries not %i" % (int(count), self._length)
jpayne@68 355 ) from None
jpayne@68 356 (self._format,) = con.execute(
jpayne@68 357 "SELECT value FROM meta_data WHERE key=?;", ("format",)
jpayne@68 358 ).fetchone()
jpayne@68 359 if fmt and fmt != self._format:
jpayne@68 360 con.close()
jpayne@68 361 raise ValueError(
jpayne@68 362 f"Index file says format {self._format}, not {fmt}"
jpayne@68 363 ) from None
jpayne@68 364 try:
jpayne@68 365 (filenames_relative_to_index,) = con.execute(
jpayne@68 366 "SELECT value FROM meta_data WHERE key=?;",
jpayne@68 367 ("filenames_relative_to_index",),
jpayne@68 368 ).fetchone()
jpayne@68 369 filenames_relative_to_index = (
jpayne@68 370 filenames_relative_to_index.upper() == "TRUE"
jpayne@68 371 )
jpayne@68 372 except TypeError:
jpayne@68 373 # Original behaviour, assume if meta_data missing
jpayne@68 374 filenames_relative_to_index = False
jpayne@68 375 self._filenames = [
jpayne@68 376 row[0]
jpayne@68 377 for row in con.execute(
jpayne@68 378 "SELECT name FROM file_data ORDER BY file_number;"
jpayne@68 379 ).fetchall()
jpayne@68 380 ]
jpayne@68 381 if filenames_relative_to_index:
jpayne@68 382 # Not implicitly relative to $PWD, explicitly relative to index file
jpayne@68 383 relative_path = os.path.abspath(os.path.dirname(index_filename))
jpayne@68 384 tmp = []
jpayne@68 385 for f in self._filenames:
jpayne@68 386 if os.path.isabs(f):
jpayne@68 387 tmp.append(f)
jpayne@68 388 else:
jpayne@68 389 # Would be stored with Unix / path separator, so convert
jpayne@68 390 # it to the local OS path separator here:
jpayne@68 391 tmp.append(
jpayne@68 392 os.path.join(relative_path, f.replace("/", os.path.sep))
jpayne@68 393 )
jpayne@68 394 self._filenames = tmp
jpayne@68 395 del tmp
jpayne@68 396 if filenames and len(filenames) != len(self._filenames):
jpayne@68 397 con.close()
jpayne@68 398 raise ValueError(
jpayne@68 399 "Index file says %i files, not %i"
jpayne@68 400 % (len(self._filenames), len(filenames))
jpayne@68 401 ) from None
jpayne@68 402 if filenames and filenames != self._filenames:
jpayne@68 403 for old, new in zip(self._filenames, filenames):
jpayne@68 404 # Want exact match (after making relative to the index above)
jpayne@68 405 if os.path.abspath(old) != os.path.abspath(new):
jpayne@68 406 con.close()
jpayne@68 407 if filenames_relative_to_index:
jpayne@68 408 raise ValueError(
jpayne@68 409 "Index file has different filenames, e.g. %r != %r"
jpayne@68 410 % (os.path.abspath(old), os.path.abspath(new))
jpayne@68 411 ) from None
jpayne@68 412 else:
jpayne@68 413 raise ValueError(
jpayne@68 414 "Index file has different filenames "
jpayne@68 415 "[This is an old index where any relative paths "
jpayne@68 416 "were relative to the original working directory]. "
jpayne@68 417 "e.g. %r != %r"
jpayne@68 418 % (os.path.abspath(old), os.path.abspath(new))
jpayne@68 419 ) from None
jpayne@68 420 # Filenames are equal (after imposing abspath)
jpayne@68 421 except sqlite3.OperationalError as err:
jpayne@68 422 con.close()
jpayne@68 423 raise ValueError(f"Not a Biopython index database? {err}") from None
jpayne@68 424 # Now we have the format (from the DB if not given to us),
jpayne@68 425 if not proxy_factory(self._format):
jpayne@68 426 con.close()
jpayne@68 427 raise ValueError(f"Unsupported format '{self._format}'")
jpayne@68 428
jpayne@68 429 def _build_index(self):
jpayne@68 430 """Call from __init__ to create a new index (PRIVATE)."""
jpayne@68 431 index_filename = self._index_filename
jpayne@68 432 relative_path = self._relative_path
jpayne@68 433 filenames = self._filenames
jpayne@68 434 fmt = self._format
jpayne@68 435 key_function = self._key_function
jpayne@68 436 proxy_factory = self._proxy_factory
jpayne@68 437 max_open = self._max_open
jpayne@68 438 random_access_proxies = self._proxies
jpayne@68 439
jpayne@68 440 if not fmt or not filenames:
jpayne@68 441 raise ValueError(
jpayne@68 442 f"Filenames to index and format required to build {index_filename!r}"
jpayne@68 443 )
jpayne@68 444 if not proxy_factory(fmt):
jpayne@68 445 raise ValueError(f"Unsupported format '{fmt}'")
jpayne@68 446 # Create the index
jpayne@68 447 con = sqlite3.dbapi2.connect(index_filename)
jpayne@68 448 self._con = con
jpayne@68 449 # print("Creating index")
jpayne@68 450 # Sqlite PRAGMA settings for speed
jpayne@68 451 con.execute("PRAGMA synchronous=OFF")
jpayne@68 452 con.execute("PRAGMA locking_mode=EXCLUSIVE")
jpayne@68 453 # Don't index the key column until the end (faster)
jpayne@68 454 # con.execute("CREATE TABLE offset_data (key TEXT PRIMARY KEY, "
jpayne@68 455 # "offset INTEGER);")
jpayne@68 456 con.execute("CREATE TABLE meta_data (key TEXT, value TEXT);")
jpayne@68 457 con.execute("INSERT INTO meta_data (key, value) VALUES (?,?);", ("count", -1))
jpayne@68 458 con.execute("INSERT INTO meta_data (key, value) VALUES (?,?);", ("format", fmt))
jpayne@68 459 con.execute(
jpayne@68 460 "INSERT INTO meta_data (key, value) VALUES (?,?);",
jpayne@68 461 ("filenames_relative_to_index", "True"),
jpayne@68 462 )
jpayne@68 463 # TODO - Record the file size and modified date?
jpayne@68 464 con.execute("CREATE TABLE file_data (file_number INTEGER, name TEXT);")
jpayne@68 465 con.execute(
jpayne@68 466 "CREATE TABLE offset_data (key TEXT, "
jpayne@68 467 "file_number INTEGER, offset INTEGER, length INTEGER);"
jpayne@68 468 )
jpayne@68 469 count = 0
jpayne@68 470 for file_index, filename in enumerate(filenames):
jpayne@68 471 # Default to storing as an absolute path,
jpayne@68 472 f = os.path.abspath(filename)
jpayne@68 473 if not os.path.isabs(filename) and not os.path.isabs(index_filename):
jpayne@68 474 # Since user gave BOTH filename & index as relative paths,
jpayne@68 475 # we will store this relative to the index file even though
jpayne@68 476 # if it may now start ../ (meaning up a level)
jpayne@68 477 # Note for cross platform use (e.g. shared drive over SAMBA),
jpayne@68 478 # convert any Windows slash into Unix style for rel paths.
jpayne@68 479 f = os.path.relpath(filename, relative_path).replace(os.path.sep, "/")
jpayne@68 480 elif (os.path.dirname(os.path.abspath(filename)) + os.path.sep).startswith(
jpayne@68 481 relative_path + os.path.sep
jpayne@68 482 ):
jpayne@68 483 # Since sequence file is in same directory or sub directory,
jpayne@68 484 # might as well make this into a relative path:
jpayne@68 485 f = os.path.relpath(filename, relative_path).replace(os.path.sep, "/")
jpayne@68 486 assert not f.startswith("../"), f
jpayne@68 487 # print("DEBUG - storing %r as [%r] %r" % (filename, relative_path, f))
jpayne@68 488 con.execute(
jpayne@68 489 "INSERT INTO file_data (file_number, name) VALUES (?,?);",
jpayne@68 490 (file_index, f),
jpayne@68 491 )
jpayne@68 492 random_access_proxy = proxy_factory(fmt, filename)
jpayne@68 493 if key_function:
jpayne@68 494 offset_iter = (
jpayne@68 495 (key_function(key), file_index, offset, length)
jpayne@68 496 for (key, offset, length) in random_access_proxy
jpayne@68 497 )
jpayne@68 498 else:
jpayne@68 499 offset_iter = (
jpayne@68 500 (key, file_index, offset, length)
jpayne@68 501 for (key, offset, length) in random_access_proxy
jpayne@68 502 )
jpayne@68 503 while True:
jpayne@68 504 batch = list(itertools.islice(offset_iter, 100))
jpayne@68 505 if not batch:
jpayne@68 506 break
jpayne@68 507 # print("Inserting batch of %i offsets, %s ... %s"
jpayne@68 508 # % (len(batch), batch[0][0], batch[-1][0]))
jpayne@68 509 con.executemany(
jpayne@68 510 "INSERT INTO offset_data (key,file_number,offset,length) VALUES (?,?,?,?);",
jpayne@68 511 batch,
jpayne@68 512 )
jpayne@68 513 con.commit()
jpayne@68 514 count += len(batch)
jpayne@68 515 if len(random_access_proxies) < max_open:
jpayne@68 516 random_access_proxies[file_index] = random_access_proxy
jpayne@68 517 else:
jpayne@68 518 random_access_proxy._handle.close()
jpayne@68 519 self._length = count
jpayne@68 520 # print("About to index %i entries" % count)
jpayne@68 521 try:
jpayne@68 522 con.execute(
jpayne@68 523 "CREATE UNIQUE INDEX IF NOT EXISTS key_index ON offset_data(key);"
jpayne@68 524 )
jpayne@68 525 except sqlite3.IntegrityError as err:
jpayne@68 526 self._proxies = random_access_proxies
jpayne@68 527 self.close()
jpayne@68 528 con.close()
jpayne@68 529 raise ValueError(f"Duplicate key? {err}") from None
jpayne@68 530 con.execute("PRAGMA locking_mode=NORMAL")
jpayne@68 531 con.execute("UPDATE meta_data SET value = ? WHERE key = ?;", (count, "count"))
jpayne@68 532 con.commit()
jpayne@68 533 # print("Index created")
jpayne@68 534
jpayne@68 535 def __repr__(self):
jpayne@68 536 return self._repr
jpayne@68 537
jpayne@68 538 def __contains__(self, key):
jpayne@68 539 return bool(
jpayne@68 540 self._con.execute(
jpayne@68 541 "SELECT key FROM offset_data WHERE key=?;", (key,)
jpayne@68 542 ).fetchone()
jpayne@68 543 )
jpayne@68 544
jpayne@68 545 def __len__(self):
jpayne@68 546 """Return the number of records indexed."""
jpayne@68 547 return self._length
jpayne@68 548 # return self._con.execute("SELECT COUNT(key) FROM offset_data;").fetchone()[0]
jpayne@68 549
jpayne@68 550 def __iter__(self):
jpayne@68 551 """Iterate over the keys."""
jpayne@68 552 for row in self._con.execute(
jpayne@68 553 "SELECT key FROM offset_data ORDER BY file_number, offset;"
jpayne@68 554 ):
jpayne@68 555 yield str(row[0])
jpayne@68 556
jpayne@68 557 def __getitem__(self, key):
jpayne@68 558 """Return record for the specified key."""
jpayne@68 559 # Pass the offset to the proxy
jpayne@68 560 row = self._con.execute(
jpayne@68 561 "SELECT file_number, offset FROM offset_data WHERE key=?;", (key,)
jpayne@68 562 ).fetchone()
jpayne@68 563 if not row:
jpayne@68 564 raise KeyError
jpayne@68 565 file_number, offset = row
jpayne@68 566 proxies = self._proxies
jpayne@68 567 if file_number in proxies:
jpayne@68 568 record = proxies[file_number].get(offset)
jpayne@68 569 else:
jpayne@68 570 if len(proxies) >= self._max_open:
jpayne@68 571 # Close an old handle...
jpayne@68 572 proxies.popitem()[1]._handle.close()
jpayne@68 573 # Open a new handle...
jpayne@68 574 proxy = self._proxy_factory(self._format, self._filenames[file_number])
jpayne@68 575 record = proxy.get(offset)
jpayne@68 576 proxies[file_number] = proxy
jpayne@68 577 if self._key_function:
jpayne@68 578 key2 = self._key_function(record.id)
jpayne@68 579 else:
jpayne@68 580 key2 = record.id
jpayne@68 581 if key != key2:
jpayne@68 582 raise ValueError(f"Key did not match ({key} vs {key2})")
jpayne@68 583 return record
jpayne@68 584
jpayne@68 585 def get_raw(self, key):
jpayne@68 586 """Return the raw record from the file as a bytes string.
jpayne@68 587
jpayne@68 588 If the key is not found, a KeyError exception is raised.
jpayne@68 589 """
jpayne@68 590 # Pass the offset to the proxy
jpayne@68 591 row = self._con.execute(
jpayne@68 592 "SELECT file_number, offset, length FROM offset_data WHERE key=?;", (key,)
jpayne@68 593 ).fetchone()
jpayne@68 594 if not row:
jpayne@68 595 raise KeyError
jpayne@68 596 file_number, offset, length = row
jpayne@68 597 proxies = self._proxies
jpayne@68 598 if file_number in proxies:
jpayne@68 599 if length:
jpayne@68 600 # Shortcut if we have the length
jpayne@68 601 h = proxies[file_number]._handle
jpayne@68 602 h.seek(offset)
jpayne@68 603 return h.read(length)
jpayne@68 604 else:
jpayne@68 605 return proxies[file_number].get_raw(offset)
jpayne@68 606 else:
jpayne@68 607 # This code is duplicated from __getitem__ to avoid a function call
jpayne@68 608 if len(proxies) >= self._max_open:
jpayne@68 609 # Close an old handle...
jpayne@68 610 proxies.popitem()[1]._handle.close()
jpayne@68 611 # Open a new handle...
jpayne@68 612 proxy = self._proxy_factory(self._format, self._filenames[file_number])
jpayne@68 613 proxies[file_number] = proxy
jpayne@68 614 if length:
jpayne@68 615 # Shortcut if we have the length
jpayne@68 616 h = proxy._handle
jpayne@68 617 h.seek(offset)
jpayne@68 618 return h.read(length)
jpayne@68 619 else:
jpayne@68 620 return proxy.get_raw(offset)
jpayne@68 621
jpayne@68 622 def close(self):
jpayne@68 623 """Close any open file handles."""
jpayne@68 624 proxies = self._proxies
jpayne@68 625 while proxies:
jpayne@68 626 proxies.popitem()[1]._handle.close()