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