p4treemap.py #6

#!/usr/bin/env python3
# -*- encoding: UTF8 -*-

"""
    P4TreeMap.py - print out tree-related information on file sizes in a
    repository or part of one.

    Data can be output in JSON hierarchy or CSV flat file.
    Main usage is as part of p4webtreemap.py which is a Flask app to display
    the results on a web page.
"""

from __future__ import print_function

import sys
import P4
import argparse
import os
import logging

DEFAULT_VERBOSITY = 'INFO'
DEFAULT_LOG_FILE = 'log-p4treemap.log'
LOGGER_NAME = 'LOG2SQL'


class P4TreeMap():

    def __init__(self, options, outstream=None):
        self.options = options
        self.options.logfile = DEFAULT_LOG_FILE
        self.p4 = P4.P4()
        if options.port:
            self.p4.port = options.port
        if options.user:
            self.p4.user = options.user
        if outstream:
            self.outstream = outstream
        else:
            self.outstream = sys.stdout
        if options.root:
            os.chdir(options.root)
        self.init_logger()
        self.logger.debug("p4treemap.py: %s" % self.options)
        self.logger.debug("Currdir: %s" % os.getcwd())
        self.p4.connect()

    def init_logger(self):
        self.logger = logging.getLogger(LOGGER_NAME)
        self.logger.setLevel(self.options.verbosity)
        outformatter = logging.Formatter('%(message)s')
        ch = logging.StreamHandler(self.outstream)
        ch.setLevel(logging.INFO)
        ch.setFormatter(outformatter)
        self.logger.addHandler(ch)
        if self.options.verbosity != logging.INFO and self.options.logfile:
            formatter = logging.Formatter('%(asctime)s:%(levelname)s %(message)s')
            fh = logging.FileHandler(self.options.logfile, mode='w')
            fh.setFormatter(formatter)
            self.logger.addHandler(fh)

    def output(self, line):
        self.logger.info(line)

    def run(self):
        seen = {}
        sep = "\t"
        self.output("id{}value".format(sep))
        for f in self.p4.run_sizes(self.options.path):
            path = f['depotFile'][2:]  # strip //
            parts = path.split('/')
            for i in range(len(parts)):
                dir = '/'.join(parts[:i])
                if dir and not dir in seen:
                    self.output("{}{}".format(dir, sep))
                    seen[dir] = 1
            self.output("{}{}{}".format(path, sep, f['fileSize']))

def main():
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument('-p', '--port', help="Perforce P4PORT to use", default=None)
    parser.add_argument('-u', '--user', help="Perforce P4USER to use", default=None)
    parser.add_argument('-r', '--root', help="Root of a DVCS server to use", default=None)
    parser.add_argument('-m', '--max', help="Maximum levels to recurse down, 0 = all", default=4)
    parser.add_argument('path', help="Perforce depot path to analyse. Default //...", default="//...")
    parser.add_argument('-v', '--verbosity', nargs='?', const="INFO", default=DEFAULT_VERBOSITY,
                        choices=('DEBUG', 'INFO', 'WARNING', 'ERROR', 'FATAL'),
                        help="Output verbosity level. Default is: " + DEFAULT_VERBOSITY)
    try:
        options = parser.parse_args()
    except Exception as e:
        parser.print_help()
        sys.exit(1)

    tree = P4TreeMap(options)
    tree.run()

if __name__ == '__main__':
    main()