1#!/usr/bin/env python3
2# Copyright 2020 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import re
7import subprocess
8
9import utils
10
11
12GIT_LSTREE_RE_LINE = re.compile(rb'^([^ ]*) ([^ ]*) ([^ ]*)\t(.*)$')
13
14
15class LazyTree:
16    """LazyTree does git mktree lazily."""
17
18    def __init__(self, treehash=None):
19        """Initializes a LazyTree.
20
21        If treehash is not None, it initializes as the tree object.
22
23        Args:
24            treehash: tree object id. please do not use a treeish, it will fail
25                later.
26        """
27        if treehash:
28            self._treehash = treehash # tree object id of current tree
29            self._subtrees = None # map from directory name to sub LazyTree
30            self._files = None # map from file naem to utils.GitFile
31            return
32        # Initialize an empty LazyTree
33        self._treehash = None
34        self._subtrees = {}
35        self._files = {}
36
37    def _loadtree(self):
38        """Loads _treehash into _subtrees and _files."""
39        if self._files is not None: # _subtrees is also not None too here.
40            return
41        output = subprocess.check_output(['git', 'ls-tree', self._treehash]).split(b'\n')
42        self._files = {}
43        self._subtrees = {}
44        for line in output:
45            if not line:
46                continue
47            m = GIT_LSTREE_RE_LINE.match(line)
48            mode, gittype, objecthash, name = m.groups()
49            assert gittype == b'blob' or gittype == b'tree'
50            assert name not in self._files and name not in self._subtrees
51            if gittype == b'blob':
52                self._files[name] = utils.GitFile(None, mode, objecthash)
53            elif gittype == b'tree':
54                self._subtrees[name] = LazyTree(objecthash)
55
56    def _remove(self, components):
57        """Removes components from self tree.
58
59        Args:
60            components: the path to remove, relative to self. Each element means
61                one level of directory tree.
62        """
63        self._loadtree()
64        self._treehash = None
65        if len(components) == 1:
66            del self._files[components[0]]
67            return
68
69        # Remove from subdirectory
70        dirname, components = components[0], components[1:]
71        subdir = self._subtrees[dirname]
72        subdir._remove(components)
73        if subdir.is_empty():
74            del self._subtrees[dirname]
75
76    def __delitem__(self, path):
77        """Removes path from self tree.
78
79        Args:
80            path: the path to remove, relative to self.
81        """
82        components = path.split(b'/')
83        self._remove(components)
84
85    def _get(self, components):
86        """Returns a file at components in utils.GitFile from self tree.
87
88        Args:
89            components: path in list instead of separated by /.
90        """
91        self._loadtree()
92        if len(components) == 1:
93            return self._files[components[0]]
94
95        dirname, components = components[0], components[1:]
96        return self._subtrees[dirname]._get(components)
97
98    def __getitem__(self, path):
99        """Returns a file at path in utils.GitFile from tree.
100
101        Args:
102            path: path of the file to read.
103        """
104        components = path.split(b'/')
105        return self._get(components)
106
107    def _set(self, components, f):
108        """Adds or replace a file.
109
110        Args:
111            components: the path to set, relative to self. Each element means
112                one level of directory tree.
113            f: a utils.GitFile object.
114        """
115
116        self._loadtree()
117        self._treehash = None
118        if len(components) == 1:
119            self._files[components[0]] = f
120            return
121
122        # Add to subdirectory
123        dirname, components = components[0], components[1:]
124        if dirname not in self._subtrees:
125            self._subtrees[dirname] = LazyTree()
126        self._subtrees[dirname]._set(components, f)
127
128    def __setitem__(self, path, f):
129        """Adds or replaces a file.
130
131        Args:
132            path: the path to set, relative to self
133            f: a utils.GitFile object
134        """
135        assert f.path.endswith(path)
136        components = path.split(b'/')
137        self._set(components, f)
138
139    def is_empty(self):
140        """Returns if self is an empty tree."""
141        return not self._subtrees and not self._files
142
143    def hash(self):
144        """Returns the hash of current tree object.
145
146        If the object doesn't exist, create it.
147        """
148        if not self._treehash:
149            self._treehash = self._mktree()
150        return self._treehash
151
152    def _mktree(self):
153        """Recreates a tree object recursively.
154
155        Lazily if subtree is unchanged.
156        """
157        keys = list(self._files.keys()) + list(self._subtrees.keys())
158        mktree_input = []
159        for name in sorted(keys):
160            file = self._files.get(name)
161            if file:
162                mktree_input.append(b'%s blob %s\t%s' % (file.mode, file.id,
163                                                         name))
164            else:
165                mktree_input.append(
166                    b'040000 tree %s\t%s' % (self._subtrees[name].hash(), name))
167        return subprocess.check_output(
168            ['git', 'mktree'],
169            input=b'\n'.join(mktree_input)).strip(b'\n')
170