<commit_before>import codecs
import hashlib
from sqlalchemy import create_engine
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import sessionmaker

from . import models
from .connect import DBConnect

class Paster():
    def __enter__(self):
        connection = DBConnect(dialect='sqlite', dbname='test.db').connect
        self.engine = create_engine(connection)
        Session = sessionmaker(bind=self.engine)
        self.session = Session()
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.session.close()
        self.engine.dispose()

    def create(self, data, ip=None, mac=None, mime=None, sunset=None):
        sha1 = hashlib.sha1(data).hexdigest()
        paste = models.Paste(
                hashid = sha1,
                ip = ip,
                mac = mac,
                mime = mime,
                sunset = sunset,
                data = data
                )
        try:
            self.session.add(paste)
            self.session.commit()
        except IntegrityError:
            paste.id = 'HASH COLLISION'
            self.session.rollback()
        return {'id': paste.id, 'hashid': sha1}

    def query(self, id=None, hashid=None):
        result = None
        if id:
            result = (self.session.query(models.Paste)
                      .filter(models.Paste.id == id).first())
        elif hashid:
            result = (self.session.query(models.Paste)
                      .filter(models.Paste.hashid == hashid).first())
        else:
            return None
        if result:
            result = {
                    'id': result.id,
                    'hashid': result.hashid,
                    'ip': result.ip,
                    'mac': result.mac,
                    'mime': result.mime,
                    'timestamp': result.timestamp,
                    'sunset': result.sunset,
                    'data': result.data
                    }

        return result

def main():
    with Paster() as paste:
        created = paste.create(b'This is a test paste')
        print(created)
        lookup = paste.query(hashid=created['hashid'])
        return lookup

if __name__ == "__main__":
    print(main())
<commit_msg>Added delete function, now only functions as a library<commit_after>import codecs
import hashlib
from sqlalchemy import create_engine
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import sessionmaker

from . import models
from .connect import DBConnect

class Paster():
    def __init__(self, dialect='sqlite', dbname='test.db'):
        self.dialect = dialect
        self.dbname = dbname

    def __enter__(self):
        connection = DBConnect(dialect=self.dialect, dbname=self.dbname).connect
        self.engine = create_engine(connection)
        Session = sessionmaker(bind=self.engine)
        self.session = Session()
        return self

    def __exit__(self, exception_type, exception_value, traceback):
        self.session.close()
        self.engine.dispose()

    def create(self, data, ip=None, mac=None, mime=None, sunset=None,
               timestamp=None):
        sha1 = hashlib.sha1(data).hexdigest()
        paste = models.Paste(
                hashid = sha1,
                ip = ip,
                mac = mac,
                mime = mime,
                sunset = sunset,
                timestamp = timestamp,
                data = data
                )
        try:
            self.session.add(paste)
            self.session.commit()
        except IntegrityError:
            paste.id = 'HASH COLLISION'
            self.session.rollback()
        return {'id': paste.id, 'hashid': sha1}

    def query(self, id=None, hashid=None):
        result = None
        if id:
            result = (self.session.query(models.Paste)
                      .filter(models.Paste.id == id).first())
        elif hashid:
            result = (self.session.query(models.Paste)
                      .filter(models.Paste.hashid == hashid).first())
        else:
            return None
        if result:
            result = {
                    'id': result.id,
                    'hashid': result.hashid,
                    'ip': result.ip,
                    'mac': result.mac,
                    'mime': result.mime,
                    'timestamp': result.timestamp,
                    'sunset': result.sunset,
                    'data': result.data
                    }

        return result

    def delete(self, id=None, hashid=None):
        if id:
            result = (self.session.query(models.Paste)
                      .filter(models.Paste.id == id).first())
        elif hashid:
            result = (self.session.query(models.Paste)
                      .filter(models.Paste.hashid == hashid).first())
        else:
            return None
        self.session.delete(result)
        self.session.commit()

{'fraction_non_alphanumeric': 0.10414971521562245, 'fraction_numerical': 0.0016273393002441008, 'mean_word_length': 1.591987348444913, 'pattern_counts': {'":': 1, '<': 3, '<?xml version=': 0, '>': 3, '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': '6831369', 'n_tokens_mistral': 1381, 'n_tokens_neox': 1308, 'n_words': 309}