spacesavers2_pdq

#!/usr/bin/env python3
# pqd = pretty darn quick

from src.VersionCheck import version_check
from src.VersionCheck import __version__
from src.utils import *

version_check()

# import required modules
import textwrap
import tqdm
import sys
from src.pdq import pdq
from multiprocessing import Pool
import argparse
from pathlib import Path
import json
import pandas as pd


def task(f):
    fd = pdq()
    fd.set(f)
    return fd

def process(fd):
    # requires global bigdict
    uid = fd.get_uid()
    if not uid in bigdict: bigdict[uid]=dict()
    inode = fd.get_inode()
    if not inode in bigdict[uid]: bigdict[uid][inode]=fd.get_size()

def main():
    elog = textwrap.dedent(
        """\
    Version:
        {}
    Example:
        > spacesavers2_pdq -f /path/to/folder -p 4 -o /path/to/output_file
        """.format(
            __version__
        )
    )
    parser = argparse.ArgumentParser(
        description="spacesavers2_pdq: get quick per user info (number of files and bytes).",
        epilog=elog,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-f",
        "--folder",
        dest="folder",
        required=True,
        type=str,
        help="spacesavers2_pdq will be run on all files in this folder and its subfolders",
    )
    parser.add_argument(
        "-p",
        "--threads",
        dest="threads",
        required=False,
        type=int,
        default=4,
        help="number of threads to be used (default 4)",
    )
    parser.add_argument(
        "-o",
        "--outfile",
        dest="outfile",
        required=False,
        type=str,
        help="outfile ...  by default output is printed to screen",
    )
    parser.add_argument(
        "-j",
        "--json",
        dest="json",
        required=False,
        type=str,
        help="outfile file in JSON format.",
    )
    parser.add_argument(
        "-q",
        "--quiet",
        dest="quiet",
        required=False,
        action=argparse.BooleanOptionalAction,
        help="Do not show progress",
    )
    parser.add_argument("-v", "--version", action="version", version=__version__)

    global args
    args = parser.parse_args()

    folder = args.folder
    p = Path(folder).absolute()
    dirs = [p]

    tqdm_disable = False
    if args.quiet: tqdm_disable = True
    # files = [p]
    # files2 = p.glob("**/*")
    # files.extend(files2)

    if args.outfile:
        outfh = open(args.outfile, 'w')
    else:
        outfh = sys.stdout

    if args.json:
        outjson = open(args.json, 'w')

    global bigdict
    bigdict=dict()

    with Pool(processes=args.threads) as pool:
        for fd in tqdm.tqdm(pool.imap_unordered(task, scantree(p,dirs)),disable=tqdm_disable):
            if not fd.is_fld(): continue # its either a file or link or directory
            process(fd)

    # now loop through dirs
    with Pool(processes=args.threads) as pool:
        for fd in tqdm.tqdm(pool.imap_unordered(task, dirs),disable=tqdm_disable):
            if not fd.is_fld(): continue # its either a file or link or directory
            process(fd)
            
    outdict=dict()
    outdict[str(p)]=dict()
    col_names =  ['uid', 'username', 'ninodes', 'nbytes', 'human_readable']
    df  = pd.DataFrame(columns = col_names)

    for uid in bigdict.keys():
        username = get_username_groupname(uid)
        outdict[str(p)][str(uid)]=dict()
        ninodes = len(bigdict[uid])
        nbytes = 0
        for inode in bigdict[uid].keys():
            nbytes += bigdict[uid][inode]
        outdict[str(p)][str(uid)]['username']=username
        outdict[str(p)][str(uid)]['ninodes']=ninodes
        outdict[str(p)][str(uid)]['nbytes']=nbytes
        my_dict = {'uid':uid,
                   'username':username,
                   'ninodes':ninodes,
                   'nbytes':nbytes,
                   'human_readable':get_human_readable_size(nbytes)}
        df.loc[len(df)] = my_dict
        # outfh.write(f"{username}\t{ninodes}\t{nbytes}\n")

    total_ninodes = df['ninodes'].sum()    
    total_nbytes = df['nbytes'].sum()
    total_humanreadable = get_human_readable_size(total_nbytes)
    my_dict = { 'uid':0,
                'username':'allusers',
                'ninodes':total_ninodes ,
                'nbytes':total_nbytes,
                'human_readable':total_humanreadable}
    df.loc[len(df)] = my_dict
    df.sort_values(by=['nbytes'],ascending=False,inplace=True)
    df['percent'] = df['nbytes'] * 100.0 / total_nbytes
    df['percent'] = df['percent'].apply(lambda x: float("{:.2f}".format(x)))
    df.to_csv(outfh,sep="\t",index=False)    
    if args.json:
        json.dump(outdict,outjson,indent=1)
        outjson.close()

    if args.outfile:
        outfh.close()

if __name__ == "__main__":
    main()