<commit_before># Copyright 2012 the rootpy developers
# distributed under the terms of the GNU General Public License
"""
This module handles conversion of ROOT's TFile and
contained TTrees into HDF5 format with PyTables
"""
from __future__ import absolute_import

import os
import sys
import tables
import warnings

from .io import root_open, utils, TemporaryFile
from . import log; log = log[__name__]
from .extern.progressbar import ProgressBar, Bar, ETA, Percentage
from .logger.util import check_tty

from root_numpy import tree2rec, RootNumpyUnconvertibleWarning


def convert(rfile, hfile, rpath='', entries=-1, userfunc=None, selection=None):

    isatty = check_tty(sys.stdout)
    if isatty:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables.openFile(filename=hfile, mode="w", title="Data")
        own_h5file = True
    own_rootfile = False
    if isinstance(rfile, basestring):
        rfile = root_open(rfile)
        own_rootfile = True

    for dirpath, dirnames, treenames in utils.walk(
            rfile, rpath, class_pattern='TTree'):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = '/' + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        if ntrees > 1:
            log.info("Will convert %i trees in this directory" % ntrees)
        else:
            log.info("Will convert 1 tree in this directory")

        for treename in treenames:

            input_tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '%s'" %
                    input_tree.GetName())
                trees = userfunc(input_tree)

                if not isinstance(trees, list):
                    trees = [trees]

            else:
                trees = [input_tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '%s' with %i entries ..." % (
                    tree.GetName(),
                    tree.GetEntries()))

                total_entries = tree.GetEntries()
                pbar = None
                if isatty and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    table = hfile.createTable(
                        group, tree.GetName(),
                        recarray, tree.GetTitle())
                    # flush data in the table
                    table.flush()
                    # flush all pending data
                    hfile.flush()
                else:
                    # read the tree in chunks
                    offset = 0
                    while offset < total_entries or offset == 0:
                        if offset > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter("ignore",
                                        RootNumpyUnconvertibleWarning)
                                recarray = tree2rec(tree,
                                        entries=entries,
                                        offset=offset,
                                        selection=selection)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(tree,
                                    entries=entries,
                                    offset=offset,
                                    selection=selection)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            table = hfile.createTable(
                                group, tree.GetName(),
                                recarray, tree.GetTitle())
                        offset += entries
                        if offset <= total_entries and pbar is not None:
                            pbar.update(offset)<commit_msg>drop object columns in root2hdf5 (not supported by PyTables)<commit_after># Copyright 2012 the rootpy developers
# distributed under the terms of the GNU General Public License
"""
This module handles conversion of ROOT's TFile and
contained TTrees into HDF5 format with PyTables
"""
from __future__ import absolute_import

import os
import sys
import tables
import warnings

from .io import root_open, utils, TemporaryFile
from . import log; log = log[__name__]
from .extern.progressbar import ProgressBar, Bar, ETA, Percentage
from .logger.util import check_tty

from root_numpy import tree2rec, RootNumpyUnconvertibleWarning


def _drop_object_col(rec, warn=True):
    # ignore columns of type `object` since PyTables does not support these
    if rec.dtype.hasobject:
        names = []
        fields = rec.dtype.fields
        for name in rec.dtype.names:
            if fields[name][0].kind != 'O':
                names.append(name)
            elif warn:
                log.warning("ignoring unsupported object branch '%s'" % name)
        return rec[names]
    return rec


def convert(rfile, hfile, rpath='', entries=-1, userfunc=None, selection=None):

    isatty = check_tty(sys.stdout)
    if isatty:
        widgets = [Percentage(), ' ', Bar(), ' ', ETA()]

    own_h5file = False
    if isinstance(hfile, basestring):
        hfile = tables.openFile(filename=hfile, mode="w", title="Data")
        own_h5file = True
    own_rootfile = False
    if isinstance(rfile, basestring):
        rfile = root_open(rfile)
        own_rootfile = True

    for dirpath, dirnames, treenames in utils.walk(
            rfile, rpath, class_pattern='TTree'):

        # skip root
        if not dirpath and not treenames:
            continue

        # skip directories w/o trees or subdirs
        if not dirnames and not treenames:
            continue

        where_group = '/' + os.path.dirname(dirpath)
        current_dir = os.path.basename(dirpath)

        if not current_dir:
            group = hfile.root
        else:
            group = hfile.createGroup(where_group, current_dir, "")

        ntrees = len(treenames)
        if ntrees > 1:
            log.info("Will convert %i trees in this directory" % ntrees)
        else:
            log.info("Will convert 1 tree in this directory")

        for treename in treenames:

            input_tree = rfile.Get(os.path.join(dirpath, treename))

            if userfunc is not None:
                tmp_file = TemporaryFile()
                # call user-defined function on tree and get output trees
                log.info("Calling user function on tree '%s'" %
                    input_tree.GetName())
                trees = userfunc(input_tree)

                if not isinstance(trees, list):
                    trees = [trees]

            else:
                trees = [input_tree]
                tmp_file = None

            for tree in trees:

                log.info("Converting tree '%s' with %i entries ..." % (
                    tree.GetName(),
                    tree.GetEntries()))

                total_entries = tree.GetEntries()
                pbar = None
                if isatty and total_entries > 0:
                    pbar = ProgressBar(widgets=widgets, maxval=total_entries)

                if entries <= 0:
                    # read the entire tree
                    if pbar is not None:
                        pbar.start()
                    recarray = tree2rec(tree, selection=selection)
                    recarray = _drop_object_col(recarray)
                    table = hfile.createTable(
                        group, tree.GetName(),
                        recarray, tree.GetTitle())
                    # flush data in the table
                    table.flush()
                    # flush all pending data
                    hfile.flush()
                else:
                    # read the tree in chunks
                    offset = 0
                    while offset < total_entries or offset == 0:
                        if offset > 0:
                            with warnings.catch_warnings():
                                warnings.simplefilter("ignore",
                                        RootNumpyUnconvertibleWarning)
                                recarray = tree2rec(tree,
                                        entries=entries,
                                        offset=offset,
                                        selection=selection)
                            recarray = _drop_object_col(recarray, warn=False)
                            table.append(recarray)
                        else:
                            recarray = tree2rec(tree,
                                    entries=entries,
                                    offset=offset,
                                    selection=selection)
                            recarray = _drop_object_col(recarray)
                            if pbar is not None:
                                # start after any output from root_numpy
                                pbar.start()
                            table = hfile.createTable(
                                group, tree.GetName(),
                                recarray, tree.GetTitle())
                        offset += entries
                        if offset <= total_entries and pbar is not None:
                            pbar.update(offset)
{'fraction_non_alphanumeric': 0.07131599684791175, 'fraction_numerical': 0.004038613081166272, 'mean_word_length': 1.3007024699750735, 'pattern_counts': {'":': 0, '<': 9, '<?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': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1027442', 'n_tokens_mistral': 2549, 'n_tokens_neox': 2370, 'n_words': 717}