# Copyright 2012 by Wibowo Arindrarto. All rights reserved.

# Revisions Copyright 2012-2015 by Peter Cock. All rights reserved.

# This code is part of the Biopython distribution and governed by its

# license. Please see the LICENSE file that should have been included

# as part of this package.

"""Common code for SearchIO tests."""

import gzip

import os

import unittest

try:

import sqlite3

except ImportError:

sqlite3 = None

from Bio import SearchIO

from Bio.SeqRecord import SeqRecord

class SearchTestBaseClass(unittest.TestCase):

def compare_attrs(self, obj_a, obj_b, attrs):

"""Compare attribute values of two objects."""

for attr in attrs:

# don't check for contained items, they are handled separately

if attr.startswith("_items"):

continue

# get attribute values from each objects

val_a = getattr(obj_a, attr)

val_b = getattr(obj_b, attr)

# special case for HSP and HSPFragment {hit,query}

# since they are seqrecords, we compare the strings only

# comparing using compare_record is too slow

if attr in ("_hit", "_query") and (val_a is not None and val_b is not None):

# compare seq directly if it's a contiguous hsp

if isinstance(val_a, SeqRecord) and isinstance(val_b, SeqRecord):

msg = f"Comparing attribute {attr}"

self.assertEqual(str(val_a.seq), str(val_b.seq), msg=msg)

elif isinstance(val_a, list) and isinstance(val_b, list):

for seq_a, seq_b in zip(val_a, val_b):

msg = f"Comparing attribute {attr}"

self.assertEqual(str(seq_a.seq), str(seq_b.seq), msg=msg)

else:

self.assertIsInstance(val_b, type(val_a))

msg = f"Comparing attribute {attr}"

self.assertEqual(val_a, val_b)

def compare_search_obj(self, obj_a, obj_b):

"""Compare attribute values of two QueryResult objects."""

# check that both qresults contain the same instance attributes

self.assertEqual(_num_difference(obj_a, obj_b), 0)

# compare qresult attributes

# if the above assertion pass, doesn't matter if we use a or be here

self.compare_attrs(obj_a, obj_b, list(obj_a.__dict__))

# compare objects recursively if it's not an HSPFragment

if not isinstance(obj_a, SearchIO.HSPFragment):

# check the number of hits contained

msg = f"comparing {obj_a!r} vs {obj_b!r}"

self.assertEqual(len(obj_a), len(obj_b), msg=msg)

for item_a, item_b in zip(obj_a, obj_b):

self.compare_search_obj(item_a, item_b)

class CheckRaw(unittest.TestCase):

"""Base class for testing index's get_raw method."""

fmt = None # define this in subclasses!

def check_raw(self, filename, id, raw, **kwargs):

"""Index filename using keyword arguments, check get_raw(id)==raw."""

idx = SearchIO.index(filename, self.fmt, **kwargs)

raw = raw.encode()

# Anticipate cases where the raw string and/or file uses different

# newline characters ~ we set everything to \n.

new = idx.get_raw(id)

self.assertIsInstance(new, bytes, f"Didn't get bytes from {self.fmt} get_raw")

self.assertEqual(raw.replace(b"\r\n", b"\n"), new.replace(b"\r\n", b"\n"))

idx.close()

# Now again, but using SQLite backend

if sqlite3:

idx = SearchIO.index_db(":memory:", filename, self.fmt, **kwargs)

new = idx.get_raw(id)

self.assertIsInstance(

new, bytes, f"Didn't get bytes from {self.fmt} get_raw"

)

self.assertEqual(raw.replace(b"\r\n", b"\n"), new.replace(b"\r\n", b"\n"))

idx.close()

if os.path.isfile(filename + ".bgz"):

# Do the tests again with the BGZF compressed file

print(f"[BONUS {filename}.bgz]")

self.check_raw(filename + ".bgz", id, raw, **kwargs)

class CheckIndex(SearchTestBaseClass):

"""Base class for testing indexing."""

def check_index(self, filename, format, **kwargs):

if filename.endswith(".bgz"):

with gzip.open(filename) as handle:

parsed = list(SearchIO.parse(handle, format, **kwargs))

else:

parsed = list(SearchIO.parse(filename, format, **kwargs))

# compare values by index

indexed = SearchIO.index(filename, format, **kwargs)

self.assertEqual(

len(parsed),

len(indexed),

"Should be %i records in %s, index says %i"

% (len(parsed), filename, len(indexed)),

)

# compare values by index_db, only if sqlite3 is present

if sqlite3 is not None:

db_indexed = SearchIO.index_db(":memory:", [filename], format, **kwargs)

self.assertEqual(

len(parsed),

len(db_indexed),

"Should be %i records in %s, index_db says %i"

% (len(parsed), filename, len(db_indexed)),

)

for qres in parsed:

idx_qres = indexed[qres.id]

# parsed and indexed qresult are different objects!

self.assertNotEqual(id(qres), id(idx_qres))

# but they should have the same attribute values

self.compare_search_obj(qres, idx_qres)

# sqlite3 comparison, only if it's present

if sqlite3 is not None:

dbidx_qres = db_indexed[qres.id]

self.assertNotEqual(id(qres), id(dbidx_qres))

self.compare_search_obj(qres, dbidx_qres)

indexed.close()

if sqlite3 is not None:

db_indexed.close()

db_indexed._con.close()

if os.path.isfile(filename + ".bgz"):

# Do the tests again with the BGZF compressed file

print(f"[BONUS {filename}.bgz]")

self.check_index(filename + ".bgz", format, **kwargs)

def _num_difference(obj_a, obj_b):

"""Return the number of instance attributes present only in one object."""

attrs_a = set(obj_a.__dict__)

attrs_b = set(obj_b.__dict__)

diff = attrs_a.symmetric_difference(attrs_b)

privates = len([x for x in diff if x.startswith("_")])

return len(diff) - privates