# Path: loom/util.py
# @contextlib.contextmanager
# def chdir(wd):
#     oldwd = os.getcwd()
#     try:
#         os.chdir(wd)
#         yield wd
#     finally:
#         os.chdir(oldwd)
# 
# def mkdir_p(dirname):
#     'like mkdir -p'
#     if not os.path.exists(dirname):
#         try:
#             os.makedirs(dirname)
#         except OSError as e:
#             if not os.path.exists(dirname):
#                 raise e
# 
# def rm_rf(path):
#     'like rm -rf'
#     if os.path.exists(path):
#         if os.path.isdir(path):
#             shutil.rmtree(path)
#         else:
#             os.remove(path)


import os
import shutil
import glob
import parsable
import loom.store
import loom.config
import loom.runner
import loom.generate
import loom.format
import loom.datasets
import loom.schema_pb2
import loom.query
import loom.preql
from distributions.io.stream import (
    open_compressed,
    json_load,
    protobuf_stream_load,
)
from loom.util import chdir, mkdir_p, rm_rf
    ])

    groups = results['samples'][0]['groups']
    assert os.listdir(groups), 'no groups were written'
    group_counts = []
    for f in os.listdir(groups):
        group_count = 0
        for _ in protobuf_stream_load(os.path.join(groups, f)):
            group_count += 1
        group_counts.append(group_count)
    print 'group_counts: {}'.format(' '.join(map(str, group_counts)))


def _load_checkpoint(step):
    message = loom.schema_pb2.Checkpoint()
    filename = checkpoint_files(step)['checkpoint']
    with open_compressed(filename) as f:
        message.ParseFromString(f.read())
    return message


@parsable.command
def load_checkpoint(name=None, period_sec=5, debug=False):
    '''
    Grab last full checkpoint for profiling, or list available datasets.
    '''
    loom.store.require(name, ['samples.0.init', 'samples.0.shuffled'])
    inputs, results = get_paths(name, 'checkpoints')

    rm_rf(results['root'])

# Path: loom/util.py
# @contextlib.contextmanager
# def chdir(wd):
#     oldwd = os.getcwd()
#     try:
#         os.chdir(wd)
#         yield wd
#     finally:
#         os.chdir(oldwd)
# 
# def mkdir_p(dirname):
#     'like mkdir -p'
#     if not os.path.exists(dirname):
#         try:
#             os.makedirs(dirname)
#         except OSError as e:
#             if not os.path.exists(dirname):
#                 raise e
# 
# def rm_rf(path):
#     'like rm -rf'
#     if os.path.exists(path):
#         if os.path.isdir(path):
#             shutil.rmtree(path)
#         else:
#             os.remove(path)


# Path: loom/benchmark.py
import os
import shutil
import glob
import parsable
import loom.store
import loom.config
import loom.runner
import loom.generate
import loom.format
import loom.datasets
import loom.schema_pb2
import loom.query
import loom.preql
from distributions.io.stream import (
    open_compressed,
    json_load,
    protobuf_stream_load,
)
from loom.util import chdir, mkdir_p, rm_rf

    ])

    groups = results['samples'][0]['groups']
    assert os.listdir(groups), 'no groups were written'
    group_counts = []
    for f in os.listdir(groups):
        group_count = 0
        for _ in protobuf_stream_load(os.path.join(groups, f)):
            group_count += 1
        group_counts.append(group_count)
    print 'group_counts: {}'.format(' '.join(map(str, group_counts)))


def _load_checkpoint(step):
    message = loom.schema_pb2.Checkpoint()
    filename = checkpoint_files(step)['checkpoint']
    with open_compressed(filename) as f:
        message.ParseFromString(f.read())
    return message


@parsable.command
def load_checkpoint(name=None, period_sec=5, debug=False):
    '''
    Grab last full checkpoint for profiling, or list available datasets.
    '''
    loom.store.require(name, ['samples.0.init', 'samples.0.shuffled'])
    inputs, results = get_paths(name, 'checkpoints')

    rm_rf(results['root'])

    mkdir_p(results['root'])