Rooba
/
hydrus
mirror of https://github.com/hydrusnetwork/hydrus.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity
hydrus/include/HydrusTagArchive.py

292 lines
9.3 KiB
Python
Raw Blame History

import os
import sqlite3
HASH_TYPE_MD5 = 0 # 16 bytes long
HASH_TYPE_SHA1 = 1 # 20 bytes long
HASH_TYPE_SHA256 = 2 # 32 bytes long
HASH_TYPE_SHA512 = 3 # 64 bytes long
# Please feel free to use this file however you wish.
# None of this is thread-safe, though, so don't try to do anything clever.
# If you want to make a new tag archive for use in hydrus, you want to do something like:
# import HydrusTagArchive
# hta = HydrusTagArchive.HydrusTagArchive( 'my_little_archive.db' )
# hta.SetHashType( HydrusTagArchive.HASH_TYPE_MD5 )
# hta.BeginBigJob()
# for ( hash, tags ) in my_complex_mappings_generator: hta.AddMappings( hash, tags )
# -or-
# for ( hash, tag ) in my_simple_mapping_generator: hta.AddMapping( hash, tag )
# hta.CommitBigJob()
# del hta
# If you are only adding a couple tags, you can exclude the BigJob stuff. It just makes millions of sequential writes more efficient.
# Also, this manages hashes as bytes, not hex, so if you have something like:
# hash = ab156e87c5d6e215ab156e87c5d6e215
# Then go hash = hash.decode( 'hex' ) before you pass it to Add/Get/Has/SetMappings
# If you have tags that are namespaced like hydrus (e.g. series:ghost in the shell), then check out:
# GetNamespaces
# DeleteNamespaces
# and
# RebuildNamespaces
# RebuildNamespaces takes namespaces_to_exclude, if you want to curate your namespaces a little better.
# If your GetNamespaces gives garbage, then just hit DeleteNamespaces. I'll be using the result of GetNamespaces to populate
# the advanced tag options widget when people sync with these archives.
# And also feel free to contact me directly at hydrus.admin@gmail.com if you need help.
class HydrusTagArchive( object ):
def __init__( self, path ):
self._path = path
if not os.path.exists( self._path ): create_db = True
else: create_db = False
self._InitDBCursor()
if create_db: self._InitDB()
self._namespaces = { namespace for ( namespace, ) in self._c.execute( 'SELECT namespace FROM namespaces;' ) }
self._namespaces.add( '' )
def _AddMappings( self, hash_id, tag_ids ):
self._c.executemany( 'INSERT OR IGNORE INTO mappings ( hash_id, tag_id ) VALUES ( ?, ? );', ( ( hash_id, tag_id ) for tag_id in tag_ids ) )
def _InitDB( self ):
self._c.execute( 'CREATE TABLE hash_type ( hash_type INTEGER );', )
self._c.execute( 'CREATE TABLE hashes ( hash_id INTEGER PRIMARY KEY, hash BLOB_BYTES );' )
self._c.execute( 'CREATE UNIQUE INDEX hashes_hash_index ON hashes ( hash );' )
self._c.execute( 'CREATE TABLE mappings ( hash_id INTEGER, tag_id INTEGER, PRIMARY KEY( hash_id, tag_id ) );' )
self._c.execute( 'CREATE INDEX mappings_hash_id_index ON mappings ( hash_id );' )
self._c.execute( 'CREATE TABLE namespaces ( namespace TEXT );' )
self._c.execute( 'CREATE TABLE tags ( tag_id INTEGER PRIMARY KEY, tag TEXT );' )
self._c.execute( 'CREATE UNIQUE INDEX tags_tag_index ON tags ( tag );' )
def _InitDBCursor( self ):
self._db = sqlite3.connect( self._path, isolation_level = None, detect_types = sqlite3.PARSE_DECLTYPES )
self._c = self._db.cursor()
def _GetHashId( self, hash, read_only = False ):
result = self._c.execute( 'SELECT hash_id FROM hashes WHERE hash = ?;', ( sqlite3.Binary( hash ), ) ).fetchone()
if result is None:
if read_only: raise Exception()
self._c.execute( 'INSERT INTO hashes ( hash ) VALUES ( ? );', ( sqlite3.Binary( hash ), ) )
hash_id = self._c.lastrowid
else: ( hash_id, ) = result
return hash_id
def _GetTagId( self, tag ):
if ':' in tag:
( namespace, subtag ) = tag.split( ':', 1 )
if namespace != '' and namespace not in self._namespaces:
self._c.execute( 'INSERT INTO namespaces ( namespace ) VALUES ( ? );', ( namespace, ) )
self._namespaces.add( namespace )
result = self._c.execute( 'SELECT tag_id FROM tags WHERE tag = ?;', ( tag, ) ).fetchone()
if result is None:
self._c.execute( 'INSERT INTO tags ( tag ) VALUES ( ? );', ( tag, ) )
tag_id = self._c.lastrowid
else: ( tag_id, ) = result
return tag_id
def BeginBigJob( self ): self._c.execute( 'BEGIN IMMEDIATE' )
def CommitBigJob( self ):
self._c.execute( 'COMMIT' )
self._c.execute( 'VACUUM' )
def AddMapping( self, hash, tag ):
hash_id = self._GetHashId( hash )
tag_id = self._GetTagId( tag )
self._c.execute( 'INSERT OR IGNORE INTO mappings ( hash_id, tag_id ) VALUES ( ?, ? );', ( hash_id, tag_id ) )
def AddMappings( self, hash, tags ):
hash_id = self._GetHashId( hash )
tag_ids = [ self._GetTagId( tag ) for tag in tags ]
self._AddMappings( hash_id, tag_ids )
def DeleteMapping( self, hash, tag ):
hash_id = self._GetHashId( hash )
tag_id = self._GetTagId( tag )
self._c.execute( 'DELETE FROM mappings WHERE hash_id = ? AND tag_id = ?;', ( hash_id, tag_id ) )
def DeleteMappings( self, hash ): self.DeleteTags( hash )
def DeleteTags( self, hash ):
try: hash_id = self._GetHashId( hash, read_only = True )
except: return
self._c.execute( 'DELET FROM mappings WHERE hash_id = ?;', ( hash_id, ) )
def DeleteNamespaces( self ):
self._namespaces = {}
self._namespaces.add( '' )
self._c.execute( 'DELETE FROM namespaces;' )
def GetHashType( self ):
try: ( hash_type, ) = self._c.execute( 'SELECT hash_type FROM hash_type;' ).fetchone()
except:
try:
( hash, ) = self._c.execute( 'SELECT hash FROM hashes;' ).fetchone()
if len( hash ) == 16: self.SetHashType( HASH_TYPE_MD5 )
elif len( hash ) == 20: self.SetHashType( HASH_TYPE_SHA1 )
elif len( hash ) == 32: self.SetHashType( HASH_TYPE_SHA256 )
elif len( hash ) == 64: self.SetHashType( HASH_TYPE_SHA512 )
else: raise Exception()
return self.GetHashType()
except: raise Exception( 'This archive has no hash type set, and as it has no files, no hash type guess can be made.' )
return hash_type
def GetMappings( self, hash ): return self.GetTags( hash )
def GetNamespaces( self ): return self._namespaces
def GetTags( self, hash ):
try: hash_id = self._GetHashId( hash, read_only = True )
except: return []
result = { tag for ( tag, ) in self._c.execute( 'SELECT tag FROM mappings, tags USING ( tag_id ) WHERE hash_id = ?;', ( hash_id, ) ) }
return result
def HasHash( self, hash ):
try:
hash_id = self._GetHashId( hash, read_only = True )
return True
except: return False
def IterateMappings( self ):
hash_ids = [ hash_id for ( hash_id, ) in self._c.execute( 'SELECT hash_id FROM hashes;' ) ]
for hash_id in hash_ids:
( hash, ) = self._c.execute( 'SELECT hash FROM hashes WHERE hash_id = ?;', ( hash_id, ) ).fetchone()
tags = self.GetTags( hash )
if len( tags ) > 0: yield ( hash, tags )
def RebuildNamespaces( self, namespaces_to_exclude = set() ):
self._namespaces = set()
self._namespaces.add( '' )
self._c.execute( 'DELETE FROM namespaces;' )
for ( tag, ) in self._c.execute( 'SELECT tag FROM tags;' ):
if ':' in tag:
( namespace, subtag ) = tag.split( ':', 1 )
if namespace != '' and namespace not in self._namespaces and namespace not in namespaces_to_exclude:
self._namespaces.add( namespace )
self._c.executemany( 'INSERT INTO namespaces ( namespace ) VALUES ( ? );', ( ( namespace, ) for namespace in self._namespaces ) )
def SetHashType( self, hash_type ):
self._c.execute( 'DELETE FROM hash_type;' )
self._c.execute( 'INSERT INTO hash_type ( hash_type ) VALUES ( ? );', ( hash_type, ) )
def SetMappings( self, hash, tags ):
hash_id = self._GetHashId( hash )
self._c.execute( 'DELETE FROM mappings WHERE hash_id = ?;', ( hash_id, ) )
tag_ids = [ self._GetTagId( tag ) for tag in tags ]
self._AddMappings( hash_id, tag_ids )
Reference in New Issue View Git Blame Copy Permalink