diff --git a/bin/autojump b/bin/autojump index 5152a0f..316f865 100755 --- a/bin/autojump +++ b/bin/autojump @@ -2,155 +2,176 @@ from __future__ import division, print_function import argparse -from sys import argv, stderr, version_info, exit, getfilesystemencoding +import sys from tempfile import NamedTemporaryFile from operator import itemgetter import os import shutil -AUTOJUMP_VERSION = "release-v20" +AUTOJUMP_VERSION = 'release-v20' MAX_KEYWEIGHT = 1000 MAX_STORED_PATHS = 600 COMPLETION_SEPARATOR = '__' -if "AUTOJUMP_DATA_DIR" in os.environ: - CONFIG_DIR = os.environ.get("AUTOJUMP_DATA_DIR") +if 'AUTOJUMP_DATA_DIR' in os.environ: + CONFIG_DIR = os.environ.get('AUTOJUMP_DATA_DIR') else: xdg_data_dir = os.environ.get('XDG_DATA_HOME') or os.path.join(os.environ['HOME'], '.local', 'share') - CONFIG_DIR=os.path.join(xdg_data_dir, 'autojump') + CONFIG_DIR = os.path.join(xdg_data_dir, 'autojump') -def uniqadd(collection, key): - """Adds a key to a list only if it is not already present""" - if key not in collection: - collection.append(key) +if CONFIG_DIR == os.path.expanduser('~'): + DB_FILE = CONFIG_DIR + '/.autojump.txt' +else: + DB_FILE = CONFIG_DIR + '/autojump.txt' -def dicadd(dic, key, increment=1): - """Increment a value in a dic, set it to 0 - if is is not already present""" - dic[key] = dic.get(key, 0.)+increment +class Database: + """ Object for interfacing with autojump database. """ + + def __init__(self, filename): + self.filename = filename + self.data = {} + self.load() + + def add(self, key, increment = 1): + """ Increment existing paths or initialize new ones to 0. """ + self.data[key] = self.data.get(key, 0.) + increment + + def decay(self): + """ Decay database entries. """ + for k in self.data.keys(): + self.data[k] *= 0.9 + + def load(self, error_recovery = False): + """ Try to open the database file, recovering from backup if needed. """ + try: + with open(self.filename, 'r') as aj_file: + for line in aj_file.readlines(): + weight, path = line[:-1].split("\t", 1) + path = decode(path, 'utf-8') + self.data[path] = float(weight) + except (IOError, EOFError): + if not error_recovery and os.path.exists(self.filename + ".bak"): + print('Problem with autojump database,\ + trying to recover from backup...', file=sys.stderr) + shutil.copy(self.filename + '.bak', self.filename) + return self.load(True) + else: + # TODO: migration code, will be removed in v22+ + # autojump_py last used in v17 + self.filename = get_db_file('autojump_py') + if os.path.exists(self.filename): + try: # fix to get optimised pickle in python < 3 + import cPickle as pickle + except ImportError: + import pickle + try: + with open(self.filename, 'rb') as aj_file: + # encoding is only specified for Python 2 compatibility + if sys.version_info[0] > 2: + self.data = pickle.load(aj_file, encoding="utf-8") + else: + self.data = pickle.load(aj_file) + unicode_dict = {} + for k, v in self.data.items(): + print(k) + unicode_dict[decode(k, errors="replace")] = v + return unicode_dict + except (IOError, EOFError, pickle.UnpicklingError): + pass + return {} # if everything fails, return an empty dictionary + + def maintenance(self): + """ Trims and decays database entries when exceeding settings. """ + if sum(self.data.values()) > MAX_KEYWEIGHT: + self.decay() + if len(self.data) > MAX_STORED_PATHS: + self.trim() + self.save() + + def save(self): + """ Save database atomically and preserve backup. """ + # check file existence and permissions + if ((not os.path.exists(self.filename)) or + os.name == 'nt' or + os.getuid() == os.stat(self.filename)[4]): + temp = NamedTemporaryFile(dir = CONFIG_DIR, delete = False) + for path, weight in sorted(self.data.items(), + key=itemgetter(1), + reverse=True): + temp.write((unico("%s\t%s\n")%(weight, path)).encode("utf-8")) + + # catching disk errors and skipping save when file handle can't be closed. + try: + # http://thunk.org/tytso/blog/2009/03/15/dont-fear-the-fsync/ + temp.flush() + os.fsync(temp) + temp.close() + except IOError as ex: + print("Error saving autojump database (disk full?)" % + ex, file=sys.stderr) + return + + shutil.move(temp.name, self.filename) + try: # backup file + import time + if (not os.path.exists(self.filename+".bak") or + time.time()-os.path.getmtime(self.filename+".bak") > 86400): + shutil.copy(self.filename, self.filename+".bak") + except OSError as ex: + print("Error while creating backup autojump file. (%s)" % + ex, file=sys.stderr) + + def trim(self): + """ If database has exceeded MAX_STORED_PATHS, removes bottom 10%. """ + dirs = list(self.data.items()) + dirs.sort(key=itemgetter(1)) + remove_cnt = .1 * MAX_STORED_PATHS + for path, _ in dirs[:remove_cnt]: + del self.data[path] + + def trim(self): + """ If database has exceeded MAX_STORED_PATHS, removes bottom 10%. """ + dirs = list(self.data.items()) + dirs.sort(key=itemgetter(1)) + remove_cnt = .1 * MAX_STORED_PATHS + for path, _ in dirs[:remove_cnt]: + del self.data[path] + + +def get_db_file(filename = "autojump.txt"): + """ Retrieve full database path. """ + # TODO: Remove when migration code is removed. + if CONFIG_DIR == os.path.expanduser("~"): + return CONFIG_DIR + "/." + filename + else: + return CONFIG_DIR + "/" + filename def output(unicode_text,encoding=None): """Wrapper for the print function, using the filesystem encoding by default to minimize encoding mismatch problems in directory names""" - if version_info[0] > 2: + if sys.version_info[0] > 2: print(unicode_text) else: if encoding is None: - encoding = getfilesystemencoding() + encoding = sys.getfilesystemencoding() print(unicode_text.encode(encoding)) def decode(text,encoding=None,errors="strict"): """Decoding step for python2.x which does not default to unicode""" - if version_info[0] > 2: + if sys.version_info[0] > 2: return text else: if encoding is None: - encoding = getfilesystemencoding() + encoding = sys.getfilesystemencoding() return text.decode(encoding,errors) def unico(text): """if python2, convert to a unicode object""" - if version_info[0] > 2: + if sys.version_info[0] > 2: return text else: return unicode(text) -def save(path_dict, dic_file): - """Save the database in an atomic way, and preserve - a backup file.""" - # If the dic_file exists and os supports permissions, check that dic_file belongs to us - # Otherwise, fail quietly - if (not os.path.exists(dic_file)) or os.name == 'nt' or os.getuid() == os.stat(dic_file)[4]: - temp = NamedTemporaryFile(dir=CONFIG_DIR, delete=False) - for path,weight in sorted(path_dict.items(),key=itemgetter(1),reverse=True): - # the db is stored in utf-8 - temp.write((unico("%s\t%s\n")%(weight,path)).encode("utf-8")) - - # Catching disk errors and skipping save since file handle can't be closed. - try: - #cf. http://thunk.org/tytso/blog/2009/03/15/dont-fear-the-fsync/ - temp.flush() - os.fsync(temp) - temp.close() - except IOError as ex: - print("Error while saving autojump database (disk full?)" % - ex, file=stderr) - return - - # Use shutil.move instead of os.rename because windows doesn't support - # using rename to overwrite files - shutil.move(temp.name, dic_file) - try: #backup file - import time - if (not os.path.exists(dic_file+".bak") or - time.time()-os.path.getmtime(dic_file+".bak")>86400): - shutil.copy(dic_file, dic_file+".bak") - except OSError as ex: - print("Error while creating backup autojump file. (%s)" % - ex, file=stderr) - -def open_dic(dic_file, error_recovery=False): - """Try hard to open the database file, recovering - from backup if needed. """ - try: - path_dict = {} - with open(dic_file, 'r') as aj_file: - for l in aj_file.readlines(): - weight,path = l[:-1].split("\t",1) - # the db is stored in utf-8 - path = decode(path,"utf-8") - path_dict[path] = float(weight) - return path_dict - except (IOError, EOFError): - if not error_recovery and os.path.exists(dic_file+".bak"): - print('Problem with autojump database,\ - trying to recover from backup...', file=stderr) - shutil.copy(dic_file+".bak", dic_file) - return open_dic(dic_file, True) - else: - # Temporary migration code - old_dic_file = get_dic_file("autojump_py") - if os.path.exists(old_dic_file): - try: # fix to get optimised pickle in python < 3 - import cPickle as pickle - except ImportError: - import pickle - try: - with open(old_dic_file, 'rb') as aj_file: - if version_info[0] > 2: - #encoding is only specified for python2.x compatibility - path_dict = pickle.load(aj_file, encoding="utf-8") - else: - path_dict = pickle.load(aj_file) - unicode_dict = {} #we now use unicode internally - for k,v in path_dict.items(): - print(k) - unicode_dict[decode(k,errors="replace")] = v - return unicode_dict - except (IOError, EOFError, pickle.UnpicklingError): - pass - return {} #if everything fails, return an empty file - - -def forget(path_dict, dic_file): - """Gradually forget about directories. Only call - from the actual jump since it can take time""" - keyweight = sum(path_dict.values()) - if keyweight > MAX_KEYWEIGHT: - for k in path_dict.keys(): - path_dict[k] *= 0.9 * MAX_KEYWEIGHT / keyweight - save(path_dict, dic_file) - -def clean_dict(sorted_dirs, path_dict): - """Limits the sized of the path_dict to MAX_STORED_PATHS. - Returns True if keys were deleted""" - if len(sorted_dirs) > MAX_STORED_PATHS: - #remove 25 more than needed, to avoid doing it every time - for path, dummy in sorted_dirs[MAX_STORED_PATHS-25:]: - del path_dict[path] - return True - else: return False - def match(path, pattern, ignore_case=False, only_end=False): """Check whether a path matches a particular pattern, and return the remaning part of the string""" @@ -185,17 +206,11 @@ def find_matches(dirs, patterns, result_list, ignore_case, max_matches, current_ #If a path doesn't exist, don't jump there #We still keep it in db in case it's from a removable drive if does_match and os.path.exists(path): - uniqadd(result_list, path) + if path not in result_list: + result_list.append(path) if len(result_list) >= max_matches : break -def get_dic_file(filename="autojump.txt"): - if CONFIG_DIR == os.path.expanduser("~"): - dic_file = CONFIG_DIR+"/." + filename - else: - dic_file = CONFIG_DIR+"/" + filename - return dic_file - def shell_utility(): """Run this when autojump is called as a shell utility""" parser = argparse.ArgumentParser(description='Automatically jump to directory passed as an argument.', @@ -214,8 +229,7 @@ def shell_utility(): help='show version information and exit') args = parser.parse_args() - dic_file = get_dic_file() - path_dict = open_dic(dic_file) + db = Database(DB_FILE) # The home dir can be reached quickly by "cd" and may interfere with other directories if (args.add): @@ -306,5 +320,4 @@ def shell_utility(): return True if __name__ == "__main__": - success=shell_utility() - if not success: exit(1) + if not shell_utility(): sys.exit(1)