fix usage info in run_benchmark

[python-delta-tar] / testing / run_benchmark.py

#!/usr/bin/env python3

# Copyright (C) 2014 Intra2net AG
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.html>

# Authors: Victor Ramirez de la Corte <victor@wadobo.com>
#          Eduardo Robles Elvira <edulix@wadobo.com>

import argparse
import os
import zlib
import fnmatch
import cProfile
import io
import pstats


def main(args = None):
    '''
    Main function, parses the command line arguments and launches the appropiate
    benchmark test.
    '''
    parser = argparse.ArgumentParser(description='Profiling test options. ')
    parser.add_argument('-t', '--test',
                        choices=['delta-tarfile', 'tarfile', 'zlib',
                                 'zlib-blocks'],
                        help='Select option for testing. ')
    parser.add_argument('-l', '--compression-level', type=int, default=9,
                        choices=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
                        help='Select level of compression for zlib. ')
    parser.add_argument('-p', '--path',
                        help='path to the file/dir to compress.')
    parser.add_argument('-P', '--profile', action='store_true',
                        help='Enable profile')
    parser.add_argument('-m', '--tarmode',  default="w:gz",
                        help='Specify tarfile mode [w:gz]')
    parser.add_argument('-st', '--strategy', default='default',
                        choices=['default', 'filtered', 'huffman_only'],
                        help='Select the deflate strategy in zlib related '
                             'tests.')
    parser.add_argument(
        '-s', '--sort',
        help='Sort output profile',
        choices=[ 'calls', 'cumulative', 'cumtime', 'file', 'path', 'module',
            'ncalls', 'pcalls', 'line', 'name', 'nfl', 'stdname', 'time',
            'tottime'],
        default="cumulative")

    if not args:
        pargs = parser.parse_args()
    else:
        pargs = parser.parse_args(args)

    compression_level = pargs.compression_level
    path = pargs.path

    # start profiling
    if pargs.profile:
        pr = cProfile.Profile()
        pr.enable()

    if pargs.test in ['delta-tarfile', 'tarfile']:
        test_tarfile(pargs.test, compression_level, path, pargs.tarmode)
    elif pargs.test == 'zlib':
        test_zlib(compression_level, path, pargs.strategy)
    elif pargs.test == 'zlib-blocks':
        test_zlib_blocks(compression_level, path, pargs.strategy)
    else:
        parser.print_help()

    # end profiling
    if pargs.profile:
        pr.disable()
        print_profile(pr, pargs.sort)

def test_tarfile(res, compression_level=9, path='source_dir', tarmode="w:gz"):
    '''
    Test that creates a tarfile called test_tarfile.tar.gz
    '''
    if res == 'delta-tarfile':
        deltatar = os.path.abspath("../deltatar")
        import sys
        sys.path.append("..")
        from deltatar import tarfile as TarFile
    elif res == 'tarfile':
        import tarfile as TarFile

    print("creating a tarfile with mode = '%s'" % tarmode)
    with open('test_tarfile.tar.gz', 'wb') as fo:
        tar = TarFile.open(mode=tarmode, fileobj=fo,
                           compresslevel=compression_level)
        tar.add(path)

def test_zlib(compression_level=9, path='source_dir', strategy='default',
              BUFSIZE = 16 * 1024):
    '''
    Test that compresses a directory with zlib
    '''
    if strategy == 'default':
        strategy = zlib.Z_DEFAULT_STRATEGY
    elif strategy ==  'filtered':
        strategy = zlib.Z_FILTERED
    else:
        strategy = zlib.Z_HUFFMAN_ONLY

    c = zlib.compressobj(compression_level,
                            zlib.DEFLATED,
                            -zlib.MAX_WBITS,
                            zlib.DEF_MEM_LEVEL,
                                strategy)

    def add(path):
        l = 0
        l_orig = 0
        if os.path.isdir(path):
            for f in os.listdir(path):
                l2, l_orig2 = add(os.path.join(path, f))
                l += l2
                l_orig += l_orig2
        elif os.path.isfile(path):
            with open(path, 'rb') as fo:
                while True:
                    buf = fo.read(BUFSIZE)
                    l_orig += len(buf)
                    if len(buf) == 0:
                        return l, l_orig
                    buf = c.compress(buf)
                    if buf is not None:
                        l += len(buf)
        return l, l_orig

    l, l_orig = add(path)
    l += len(c.flush())
    print("total size: %d MiB compressed, %d MiB uncompressed" % (
        (l / (1024 ** 2)),
        (l_orig / (1024 ** 2))))


def test_zlib_blocks(compression_level=9, path='source_dir', strategy='default',
                     BUFSIZE = 16 * 1024):
    '''
    Test that compresses a directory with zlib
    '''
    if strategy == 'default':
        strategy = zlib.Z_DEFAULT_STRATEGY
    elif strategy ==  'filtered':
        strategy = zlib.Z_FILTERED
    else:
        strategy = zlib.Z_HUFFMAN_ONLY

    def add(path):
        l = 0
        l_orig = 0
        if os.path.isdir(path):
            for f in os.listdir(path):
                l2, l_orig2 = add(os.path.join(path, f))
                l += l2
                l_orig += l_orig2
        elif os.path.isfile(path):
            c = zlib.compressobj(compression_level,
                                zlib.DEFLATED,
                                -zlib.MAX_WBITS,
                                zlib.DEF_MEM_LEVEL,
                                strategy)
            with open(path, 'rb') as fo:
                while True:
                    buf = fo.read(BUFSIZE)
                    l_orig += len(buf)
                    if len(buf) == 0:
                        l += len(c.flush())
                        return l, l_orig
                    buf = c.compress(buf)
                    if buf is not None:
                        l += len(buf)
        return l, l_orig

    l, l_orig = add(path)
    print("total size: %d MiB compressed, %d MiB uncompressed" % (
        (l / (1024 ** 2)),
        (l_orig / (1024 ** 2))))

def print_profile(pr, sort):
    s = io.StringIO()
    if not sort:
        sort = 'cumulative'
    ps = pstats.Stats(pr, stream=s).sort_stats(sort)
    ps.print_stats(20)
    print(s.getvalue())


if __name__ == '__main__':
    main()