1"""Tools for caching xonsh code.""" 2import os 3import sys 4import hashlib 5import marshal 6import builtins 7 8from xonsh import __version__ as XONSH_VERSION 9from xonsh.lazyasd import lazyobject 10from xonsh.platform import PYTHON_VERSION_INFO_BYTES 11 12 13def _splitpath(path, sofar=[]): 14 folder, path = os.path.split(path) 15 if path == "": 16 return sofar[::-1] 17 elif folder == "": 18 return (sofar + [path])[::-1] 19 else: 20 return _splitpath(folder, sofar + [path]) 21 22 23@lazyobject 24def _CHARACTER_MAP(): 25 cmap = {chr(o): "_%s" % chr(o + 32) for o in range(65, 91)} 26 cmap.update({".": "_.", "_": "__"}) 27 return cmap 28 29 30def _cache_renamer(path, code=False): 31 if not code: 32 path = os.path.realpath(path) 33 o = ["".join(_CHARACTER_MAP.get(i, i) for i in w) for w in _splitpath(path)] 34 o[-1] = "{}.{}".format(o[-1], sys.implementation.cache_tag) 35 return o 36 37 38def _make_if_not_exists(dirname): 39 if not os.path.isdir(dirname): 40 os.makedirs(dirname) 41 42 43def should_use_cache(execer, mode): 44 """ 45 Return ``True`` if caching has been enabled for this mode (through command 46 line flags or environment variables) 47 """ 48 if mode == "exec": 49 return (execer.scriptcache or execer.cacheall) and ( 50 builtins.__xonsh_env__["XONSH_CACHE_SCRIPTS"] 51 or builtins.__xonsh_env__["XONSH_CACHE_EVERYTHING"] 52 ) 53 else: 54 return execer.cacheall or builtins.__xonsh_env__["XONSH_CACHE_EVERYTHING"] 55 56 57def run_compiled_code(code, glb, loc, mode): 58 """ 59 Helper to run code in a given mode and context 60 """ 61 if code is None: 62 return 63 if mode in {"exec", "single"}: 64 func = exec 65 else: 66 func = eval 67 func(code, glb, loc) 68 69 70def get_cache_filename(fname, code=True): 71 """ 72 Return the filename of the cache for the given filename. 73 74 Cache filenames are similar to those used by the Mercurial DVCS for its 75 internal store. 76 77 The ``code`` switch should be true if we should use the code store rather 78 than the script store. 79 """ 80 datadir = builtins.__xonsh_env__["XONSH_DATA_DIR"] 81 cachedir = os.path.join( 82 datadir, "xonsh_code_cache" if code else "xonsh_script_cache" 83 ) 84 cachefname = os.path.join(cachedir, *_cache_renamer(fname, code=code)) 85 return cachefname 86 87 88def update_cache(ccode, cache_file_name): 89 """ 90 Update the cache at ``cache_file_name`` to contain the compiled code 91 represented by ``ccode``. 92 """ 93 if cache_file_name is not None: 94 _make_if_not_exists(os.path.dirname(cache_file_name)) 95 with open(cache_file_name, "wb") as cfile: 96 cfile.write(XONSH_VERSION.encode() + b"\n") 97 cfile.write(bytes(PYTHON_VERSION_INFO_BYTES) + b"\n") 98 marshal.dump(ccode, cfile) 99 100 101def _check_cache_versions(cfile): 102 # version data should be < 1 kb 103 ver = cfile.readline(1024).strip() 104 if ver != XONSH_VERSION.encode(): 105 return False 106 ver = cfile.readline(1024).strip() 107 return ver == PYTHON_VERSION_INFO_BYTES 108 109 110def compile_code(filename, code, execer, glb, loc, mode): 111 """ 112 Wrapper for ``execer.compile`` to compile the given code 113 """ 114 try: 115 if not code.endswith("\n"): 116 code += "\n" 117 old_filename = execer.filename 118 execer.filename = filename 119 ccode = execer.compile(code, glbs=glb, locs=loc, mode=mode, filename=filename) 120 except Exception: 121 raise 122 finally: 123 execer.filename = old_filename 124 return ccode 125 126 127def script_cache_check(filename, cachefname): 128 """ 129 Check whether the script cache for a particular file is valid. 130 131 Returns a tuple containing: a boolean representing whether the cached code 132 should be used, and the cached code (or ``None`` if the cache should not be 133 used). 134 """ 135 ccode = None 136 run_cached = False 137 if os.path.isfile(cachefname): 138 if os.stat(cachefname).st_mtime >= os.stat(filename).st_mtime: 139 with open(cachefname, "rb") as cfile: 140 if not _check_cache_versions(cfile): 141 return False, None 142 ccode = marshal.load(cfile) 143 run_cached = True 144 return run_cached, ccode 145 146 147def run_script_with_cache(filename, execer, glb=None, loc=None, mode="exec"): 148 """ 149 Run a script, using a cached version if it exists (and the source has not 150 changed), and updating the cache as necessary. 151 """ 152 run_cached = False 153 use_cache = should_use_cache(execer, mode) 154 cachefname = get_cache_filename(filename, code=False) 155 if use_cache: 156 run_cached, ccode = script_cache_check(filename, cachefname) 157 if not run_cached: 158 with open(filename, "r") as f: 159 code = f.read() 160 ccode = compile_code(filename, code, execer, glb, loc, mode) 161 update_cache(ccode, cachefname) 162 run_compiled_code(ccode, glb, loc, mode) 163 164 165def code_cache_name(code): 166 """ 167 Return an appropriate spoofed filename for the given code. 168 """ 169 if isinstance(code, str): 170 _code = code.encode() 171 else: 172 _code = code 173 return hashlib.md5(_code).hexdigest() 174 175 176def code_cache_check(cachefname): 177 """ 178 Check whether the code cache for a particular piece of code is valid. 179 180 Returns a tuple containing: a boolean representing whether the cached code 181 should be used, and the cached code (or ``None`` if the cache should not be 182 used). 183 """ 184 ccode = None 185 run_cached = False 186 if os.path.isfile(cachefname): 187 with open(cachefname, "rb") as cfile: 188 if not _check_cache_versions(cfile): 189 return False, None 190 ccode = marshal.load(cfile) 191 run_cached = True 192 return run_cached, ccode 193 194 195def run_code_with_cache(code, execer, glb=None, loc=None, mode="exec"): 196 """ 197 Run a piece of code, using a cached version if it exists, and updating the 198 cache as necessary. 199 """ 200 use_cache = should_use_cache(execer, mode) 201 filename = code_cache_name(code) 202 cachefname = get_cache_filename(filename, code=True) 203 run_cached = False 204 if use_cache: 205 run_cached, ccode = code_cache_check(cachefname) 206 if not run_cached: 207 ccode = compile_code(filename, code, execer, glb, loc, mode) 208 update_cache(ccode, cachefname) 209 run_compiled_code(ccode, glb, loc, mode) 210