<commit_before>from __future__ import with_statement
from datetime import datetime
from getpass import getuser
import msgpack
import os
import socket
import stat
import sys
from xattr import xattr, XATTR_NOFOLLOW

from . import NS_ARCHIVE_METADATA, NS_ARCHIVE_ITEMS, NS_ARCHIVE_CHUNKS, NS_CHUNK, \
    PACKET_ARCHIVE_METADATA, PACKET_ARCHIVE_ITEMS, PACKET_ARCHIVE_CHUNKS, PACKET_CHUNK
from ._speedups import chunkify
from .helpers import uid2user, user2uid, gid2group, group2gid, IntegrityError

CHUNK_SIZE = 64 * 1024
WINDOW_SIZE = 4096

have_lchmod = hasattr(os, 'lchmod')
linux = sys.platform == 'linux2'


class Archive(object):

    class DoesNotExist(Exception):
        pass

    def __init__(self, store, keychain, name=None):
        self.keychain = keychain
        self.store = store
        self.items = []
        self.items_ids = []
        self.hard_links = {}
        if name:
            self.load(self.keychain.id_hash(name))

    def load(self, id):
        self.id = id
        try:
            kind, data, self.hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_METADATA, self.id))
        except self.store.DoesNotExist:
            raise self.DoesNotExist
        assert kind == PACKET_ARCHIVE_METADATA
        self.metadata = msgpack.unpackb(data)
        assert self.metadata['version'] == 1

    @property
    def ts(self):
        """Timestamp of archive creation in UTC"""
        return datetime.strptime(self.metadata['time'], '%Y-%m-%dT%H:%M:%S.%f')

    def get_chunks(self):
        for id in self.metadata['chunks_ids']:
            magic, data, hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_CHUNKS, id))
            assert magic == PACKET_ARCHIVE_CHUNKS
            assert hash == id
            chunks = msgpack.unpackb(data)
            for chunk in chunks:
                yield chunk

    def get_items(self):
        for id in self.metadata['items_ids']:
            magic, data, items_hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_ITEMS, id))
            assert magic == PACKET_ARCHIVE_ITEMS
            assert items_hash == id
            items = msgpack.unpackb(data)
            for item in items:
                yield item

    def add_item(self, item):
        self.items.append(item)
        if len(self.items) > 100000:
            self.flush_items()

    def flush_items(self):
        data, hash = self.keychain.encrypt(PACKET_ARCHIVE_ITEMS, msgpack.packb(self.items))
        self.store.put(NS_ARCHIVE_ITEMS, hash, data)
        self.items_ids.append(hash)
        self.items = []

    def save_chunks(self, cache):
        chunks = []
        ids = []
        def flush(chunks):
            data, hash = self.keychain.encrypt(PACKET_ARCHIVE_CHUNKS, msgpack.packb(chunks))
            self.store.put(NS_ARCHIVE_CHUNKS, hash, data)
            ids.append(hash)
        for id, (count, size) in cache.chunks.iteritems():
            if count > 1000000:
                chunks.append((id, size))
            if len(chunks) > 100000:
                flush(chunks)
                chunks = []
        flush(chunks)
        return ids

    def save(self, name, cache):
        self.id = self.keychain.id_hash(name)
        chunks_ids = self.save_chunks(cache)
        self.flush_items()
        metadata = {
            'version': 1,
            'name': name,
            'chunks_ids': chunks_ids,
            'items_ids': self.items_ids,
            'cmdline': sys.argv,
            'hostname': socket.gethostname(),
            'username': getuser(),
            'time': datetime.utcnow().isoformat(),
        }
        data, self.hash = self.keychain.encrypt(PACKET_ARCHIVE_METADATA, msgpack.packb(metadata))
        self.store.put(NS_ARCHIVE_METADATA, self.id, data)
        self.store.commit()
        cache.commit()

    def stats(self, cache):
        osize = csize = usize = 0
        for item in self.get_items():
            if stat.S_ISREG(item['mode']) and not 'source' in item:
                osize += item['size']
        for id, size in self.get_chunks():
            csize += size
            if cache.seen_chunk(id) == 1:
                usize += size
        return osize, csize, usize

    def extract_item(self, item, dest=None):
        dest = dest or os.getcwdu()
        dir_stat_queue = []
        assert item['path'][0] not in ('/', '\\', ':')
        path = os.path.join(dest, item['path'].decode('utf-8'))
        mode = item['mode']
        if stat.S_ISDIR(mode):
            if not os.path.exists(path):
                os.makedirs(path)
            self.restore_attrs(path, item)
        elif stat.S_ISFIFO(mode):
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            os.mkfifo(path)
            self.restore_attrs(path, item)
        elif stat.S_ISLNK(mode):
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            source = item['source']
            if os.path.exists(path):
                os.unlink(path)
            os.symlink(source, path)
            self.restore_attrs(path, item, symlink=True)
        elif stat.S_ISREG(mode):
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            # Hard link?
            if 'source' in item:
                source = os.path.join(dest, item['source'])
                if os.path.exists(path):
                    os.unlink(path)
                os.link(source, path)
            else:
                with open(path, 'wb') as fd:
                    for id in item['chunks']:
                        try:
                            magic, data, hash = self.keychain.decrypt(self.store.get(NS_CHUNK, id))
                            assert magic == PACKET_CHUNK
                            if self.keychain.id_hash(data) != id:
                                raise IntegrityError('chunk hash did not match')
                            fd.write(data)
                        except ValueError:
                            raise Exception('Invalid chunk checksum')<commit_msg>Decode hardlink source path before linking<commit_after>from __future__ import with_statement
from datetime import datetime
from getpass import getuser
import msgpack
import os
import socket
import stat
import sys
from xattr import xattr, XATTR_NOFOLLOW

from . import NS_ARCHIVE_METADATA, NS_ARCHIVE_ITEMS, NS_ARCHIVE_CHUNKS, NS_CHUNK, \
    PACKET_ARCHIVE_METADATA, PACKET_ARCHIVE_ITEMS, PACKET_ARCHIVE_CHUNKS, PACKET_CHUNK
from ._speedups import chunkify
from .helpers import uid2user, user2uid, gid2group, group2gid, IntegrityError

CHUNK_SIZE = 64 * 1024
WINDOW_SIZE = 4096

have_lchmod = hasattr(os, 'lchmod')
linux = sys.platform == 'linux2'


class Archive(object):

    class DoesNotExist(Exception):
        pass

    def __init__(self, store, keychain, name=None):
        self.keychain = keychain
        self.store = store
        self.items = []
        self.items_ids = []
        self.hard_links = {}
        if name:
            self.load(self.keychain.id_hash(name))

    def load(self, id):
        self.id = id
        try:
            kind, data, self.hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_METADATA, self.id))
        except self.store.DoesNotExist:
            raise self.DoesNotExist
        assert kind == PACKET_ARCHIVE_METADATA
        self.metadata = msgpack.unpackb(data)
        assert self.metadata['version'] == 1

    @property
    def ts(self):
        """Timestamp of archive creation in UTC"""
        return datetime.strptime(self.metadata['time'], '%Y-%m-%dT%H:%M:%S.%f')

    def get_chunks(self):
        for id in self.metadata['chunks_ids']:
            magic, data, hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_CHUNKS, id))
            assert magic == PACKET_ARCHIVE_CHUNKS
            assert hash == id
            chunks = msgpack.unpackb(data)
            for chunk in chunks:
                yield chunk

    def get_items(self):
        for id in self.metadata['items_ids']:
            magic, data, items_hash = self.keychain.decrypt(self.store.get(NS_ARCHIVE_ITEMS, id))
            assert magic == PACKET_ARCHIVE_ITEMS
            assert items_hash == id
            items = msgpack.unpackb(data)
            for item in items:
                yield item

    def add_item(self, item):
        self.items.append(item)
        if len(self.items) > 100000:
            self.flush_items()

    def flush_items(self):
        data, hash = self.keychain.encrypt(PACKET_ARCHIVE_ITEMS, msgpack.packb(self.items))
        self.store.put(NS_ARCHIVE_ITEMS, hash, data)
        self.items_ids.append(hash)
        self.items = []

    def save_chunks(self, cache):
        chunks = []
        ids = []
        def flush(chunks):
            data, hash = self.keychain.encrypt(PACKET_ARCHIVE_CHUNKS, msgpack.packb(chunks))
            self.store.put(NS_ARCHIVE_CHUNKS, hash, data)
            ids.append(hash)
        for id, (count, size) in cache.chunks.iteritems():
            if count > 1000000:
                chunks.append((id, size))
            if len(chunks) > 100000:
                flush(chunks)
                chunks = []
        flush(chunks)
        return ids

    def save(self, name, cache):
        self.id = self.keychain.id_hash(name)
        chunks_ids = self.save_chunks(cache)
        self.flush_items()
        metadata = {
            'version': 1,
            'name': name,
            'chunks_ids': chunks_ids,
            'items_ids': self.items_ids,
            'cmdline': sys.argv,
            'hostname': socket.gethostname(),
            'username': getuser(),
            'time': datetime.utcnow().isoformat(),
        }
        data, self.hash = self.keychain.encrypt(PACKET_ARCHIVE_METADATA, msgpack.packb(metadata))
        self.store.put(NS_ARCHIVE_METADATA, self.id, data)
        self.store.commit()
        cache.commit()

    def stats(self, cache):
        osize = csize = usize = 0
        for item in self.get_items():
            if stat.S_ISREG(item['mode']) and not 'source' in item:
                osize += item['size']
        for id, size in self.get_chunks():
            csize += size
            if cache.seen_chunk(id) == 1:
                usize += size
        return osize, csize, usize

    def extract_item(self, item, dest=None):
        dest = dest or os.getcwdu()
        dir_stat_queue = []
        assert item['path'][0] not in ('/', '\\', ':')
        path = os.path.join(dest, item['path'].decode('utf-8'))
        mode = item['mode']
        if stat.S_ISDIR(mode):
            if not os.path.exists(path):
                os.makedirs(path)
            self.restore_attrs(path, item)
        elif stat.S_ISFIFO(mode):
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            os.mkfifo(path)
            self.restore_attrs(path, item)
        elif stat.S_ISLNK(mode):
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            source = item['source']
            if os.path.exists(path):
                os.unlink(path)
            os.symlink(source, path)
            self.restore_attrs(path, item, symlink=True)
        elif stat.S_ISREG(mode):
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            # Hard link?
            if 'source' in item:
                source = os.path.join(dest, item['source'].decode('utf-8'))
                if os.path.exists(path):
                    os.unlink(path)
                os.link(source, path)
            else:
                with open(path, 'wb') as fd:
                    for id in item['chunks']:
                        try:
                            magic, data, hash = self.keychain.decrypt(self.store.get(NS_CHUNK, id))
                            assert magic == PACKET_CHUNK
                            if self.keychain.id_hash(data) != id:
                                raise IntegrityError('chunk hash did not match')
                            fd.write(data)
                        except ValueError:
                            raise Exception('Invalid chunk checksum')
{'fraction_non_alphanumeric': 0.11041446136308118, 'fraction_numerical': 0.006595554107971664, 'mean_word_length': 1.872310570626754, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6047101', 'n_tokens_mistral': 3852, 'n_tokens_neox': 3734, 'n_words': 794}