<commit_before>from . import api
import contextlib
import gzip
import os
import sqlite3


class SqliteCorpus(api.Corpus):
    def __init__(self, *args, **kwargs):
        api.Corpus.__init__(self, *args, **kwargs)
        self._index = os.path.join(self.basedir, 'index.sqlite3')

    @contextlib.contextmanager
    def _dbcon(self):
        dbcon = sqlite3.connect(self._index)
        try:
            dbcon.row_factory = sqlite3.Row
            yield dbcon
        finally:
            dbcon.close()

    def _build_index(self):
        with self._dbcon() as dbcon:
            dbcon.execute('''
                CREATE TABLE IF NOT EXISTS TextInfo(
                    uid INTEGER PRIMARY KEY,
                    title TEXT,
                    author TEXT,
                    location TEXT
                )''')
            dbcon.executemany('''
                INSERT INTO TextInfo(uid, title, author, location)
                VALUES(?, ?, ?, ?)
            ''', ((text_info.uid, text_info.title, text_info.author,
                   os.path.join(self.basedir, '%s.gz' % text_info.uid))
                  for text_info in iter(self.text_source)))

    def _fulltext(self, text_info, location=None):
        if location is None:
            with self._dbcon() as dbcon:
                result = dbcon.execute('''
                    SELECT location
                    FROM TextInfo
                    WHERE uid = ?
                    LIMIT 1
                ''', (text_info.uid, )).fetchone()
            location = result['location']

        try:
            with gzip.open(location, 'rb') as gzipped:
                fulltext = gzipped.read()
        except IOError:
            fulltext = self.text_source.fulltext(text_info)
            if fulltext:
                with gzip.open(location, 'wb') as gzipped:
                    gzipped.write(fulltext)
        return fulltext

    def texts_for_author(self, author):
        with self._dbcon() as dbcon:
            results = dbcon.execute('''
                SELECT uid, author, title, location
                FROM TextInfo
                WHERE author LIKE ?
                ORDER BY author
            ''', ('%' + author + '%', ))

        for result in results:
            text_info = api.TextInfo(
                uid=result['uid'],
                author=result['author'],
                title=result['title'])
            yield text_info, self._fulltext(text_info, result['location'])
<commit_msg>Log warning on multiple matches for author-query<commit_after>from . import api
import collections
import contextlib
import gzip
import logging
import os
import sqlite3


class SqliteCorpus(api.Corpus):
    def __init__(self, *args, **kwargs):
        api.Corpus.__init__(self, *args, **kwargs)
        self._index = os.path.join(self.basedir, 'index.sqlite3')

    @contextlib.contextmanager
    def _dbcon(self):
        dbcon = sqlite3.connect(self._index)
        try:
            dbcon.row_factory = sqlite3.Row
            yield dbcon
        finally:
            dbcon.close()

    def _build_index(self):
        with self._dbcon() as dbcon:
            dbcon.execute('''
                CREATE TABLE IF NOT EXISTS TextInfo(
                    uid INTEGER PRIMARY KEY,
                    title TEXT,
                    author TEXT,
                    location TEXT
                )''')
            dbcon.executemany('''
                INSERT INTO TextInfo(uid, title, author, location)
                VALUES(?, ?, ?, ?)
            ''', ((text_info.uid, text_info.title, text_info.author,
                   os.path.join(self.basedir, '%s.gz' % text_info.uid))
                  for text_info in iter(self.text_source)))

    def _fulltext(self, text_info, location=None):
        if location is None:
            with self._dbcon() as dbcon:
                result = dbcon.execute('''
                    SELECT location
                    FROM TextInfo
                    WHERE uid = ?
                    LIMIT 1
                ''', (text_info.uid, )).fetchone()
            location = result['location']

        try:
            with gzip.open(location, 'rb') as gzipped:
                fulltext = gzipped.read()
        except IOError:
            fulltext = self.text_source.fulltext(text_info)
            if fulltext:
                with gzip.open(location, 'wb') as gzipped:
                    gzipped.write(fulltext)
        return fulltext

    def texts_for_author(self, author):
        matched_authors = collections.defaultdict(list)
        with self._dbcon() as dbcon:
            for row in dbcon.execute('''
                SELECT uid, author, title, location
                FROM TextInfo
                WHERE author LIKE ?
                ORDER BY author
            ''', ('%' + author + '%', )):
                matched_authors[row['author']].append(row)
        if len(matched_authors) > 1:
            logging.warning(
                '%d authors match the query "%s": %s',
                len(matched_authors), author,
                ', '.join('"%s"' % author for author in matched_authors))

        for matched_author, rows in matched_authors.iteritems():
            for row in rows:
                text_info = api.TextInfo(
                    uid=row['uid'],
                    author=matched_author,
                    title=row['title'])
                yield text_info, self._fulltext(text_info, row['location'])

{'fraction_non_alphanumeric': 0.10192837465564739, 'fraction_numerical': 0.00202020202020202, 'mean_word_length': 1.524802967083913, 'pattern_counts': {'":': 1, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5352869', 'n_tokens_mistral': 1535, 'n_tokens_neox': 1426, 'n_words': 359}