pygit2 auto treebuilder

auto_treebuilder_test.py

"""
    Defines a function `auto_insert` to help with
    `pygit2.Repository.TreeBuilder`s.
    
    Just create the top-level `TreeBuilder`, and it will handle all subtree
    creation if you give it paths.
"""

import shutil
from tempfile import mkdtemp
from pygit2 import init_repository
from pygit2 import GIT_FILEMODE_TREE, GIT_FILEMODE_BLOB

repo_dir = mkdtemp()
repo = init_repository(repo_dir)


def auto_insert(repo, treebuilder, path, thing, mode):
    """figure out and deal with the necessary subtree structure"""
    path_parts = path.split('/', 1)
    if len(path_parts) == 1:  # base case
        treebuilder.insert(path, thing, mode)
        return treebuilder.write()

    subtree_name, sub_path = path_parts
    tree_oid = treebuilder.write()
    tree = repo.get(tree_oid)
    try:
        entry = tree[subtree_name]
        assert entry.filemode == GIT_FILEMODE_TREE,\
            '{} already exists as a blob, not a tree'.format(entry.name)
        existing_subtree = repo.get(entry.hex)
        sub_treebuilder = repo.TreeBuilder(existing_subtree)
    except KeyError:
        sub_treebuilder = repo.TreeBuilder()

    subtree_oid = auto_insert(repo, sub_treebuilder, sub_path, thing, mode)
    treebuilder.insert(subtree_name, subtree_oid, GIT_FILEMODE_TREE)
    return treebuilder.write()


# create the only treebuilder we want to have to deal with
root = repo.TreeBuilder()


# make something for the tree
path = 'some-folder/filename.txt'
blob_sha = repo.create_blob('asdf')
# throw it in
auto_insert(repo, root, path, blob_sha, GIT_FILEMODE_BLOB)


# add something more to the same subtree
path2 = 'some-folder/another-file.txt'
blob2_sha = repo.create_blob('fdsa')
auto_insert(repo, root, path2, blob2_sha, GIT_FILEMODE_BLOB)

# add something deep
path3 = 'very/deeply/nested/deep-file.txt'
blob3_sha = repo.create_blob('sup')
auto_insert(repo, root, path3, blob3_sha, GIT_FILEMODE_BLOB)

# add something to the top
path4 = 'top-file.txt'
blob4_sha = repo.create_blob('blah')
auto_insert(repo, root, path4, blob4_sha, GIT_FILEMODE_BLOB)


# did it stick?
def print_tree(tree, indent=''):
    for entry in tree:
        if entry.filemode == GIT_FILEMODE_TREE:
            print('{}{}/'.format(indent, entry.name))
            subtree = repo.get(entry.hex)
            print_tree(subtree, indent='    {}'.format(indent))
        else:
            print('{}{}'.format(indent, entry.name))
tree_oid = root.write()
tree = repo.get(tree_oid)
print_tree(tree)


# housekeeping: clean up the temp repo
shutil.rmtree(repo_dir)

sample output

$ python auto_treebuilder_test.py
some-folder/
    another-file.txt
    filename.txt
top-file.txt
very/
    deeply/
        nested/
            deep-file.txt