the unbalancer.py for mergerfs

unbalancer.py

#!/usr/bin/env python2                                                                                                                                                                                                                                                        
import argparse                                                                                                                                                                                                                                                               
import os                                                                                                                                                                                                                                                                     
import shutil                                                                                                                                                                                                                                                                 
import sys                                                                                                                                                                                                                                                                    
from collections import defaultdict                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                              
def find(disks):                                                                                                                                                                                                                                                              
    cache = defaultdict(list)                                                                                                                                                                                                                                                 
    for disk in disks:                                                                                                                                                                                                                                                        
        for root, _, files in os.walk(disk):                                                                                                                                                                                                                                  
            for file_ in files:                                                                                                                                                                                                                                               
                full_path = os.path.join(root, file_)                                                                                                                                                                                                                         
                rel_path = os.path.relpath(full_path, disk)                                                                                                                                                                                                                   
                if "/" not in rel_path:                                                                                                                                                                                                                                       
                    continue                                                                                                                                                                                                                                                  
                top_dir = rel_path.split("/")[0]                                                                                                                                                                                                                              
                cache[top_dir].append(full_path)                                                                                                                                                                                                                              
    return cache                                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                              
def prefix_with_most_files(key, values):                                                                                                                                                                                                                                      
    prefixes = set(map(lambda path: path.split(key, 1)[0], values))                                                                                                                                                                                                           
                                                                                                                                                                                                                                                                              
    return sorted(prefixes,                                                                                                                                                                                                                                                   
            key=lambda prefix:                                                                                                                                                                                                                                                
                sum(map(lambda f: os.path.getsize(f) if f.startswith(prefix) else 0, values)))[-1]                                                                                                                                                                            
                                                                                                                                                                                                                                                                              
def decide(files):                                                                                                                                                                                                                                                            
    moves = []                                                                                                                                                                                                                                                                
    for key, values in files.items():                                                                                                                                                                                                                                         
        dst_prefix = prefix_with_most_files(key, values)                                                                                                                                                                                                                      
        for file_ in values:                                                                                                                                                                                                                                                  
            if file_.startswith(dst_prefix):                                                                                                                                                                                                                                  
                continue                                                                                                                                                                                                                                                      
            old_prefix = file_.split(key, 1)[0]                                                                                                                                                                                                                               
            new_path = os.path.join(dst_prefix, os.path.relpath(file_, old_prefix))                                                                                                                                                                                           
            moves.append((file_, new_path))                                                                                                                                                                                                                                   
    return moves                                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                              
def fake_move(moves):                                                                                                                                                                                                                                                         
    for src, dst in moves:                                                                                                                                                                                                                                                    
        print("[DryRun] Moving %s => %s" % (src, dst))                                                                                                                                                                                                                        
                                                                                                                                                                                                                                                                              
def move(moves):                                                                                                                                                                                                                                                              
    for src, dst in moves:                                                                                                                                                                                                                                                    
        print("Moving %s => %s" % (src, dst))                                                                                                                                                                                                                                 
        dirname = os.path.dirname(dst)                                                                                                                                                                                                                                        
        try:                                                                                                                                                                                                                                                                  
            os.makedirs(dirname)                                                                                                                                                                                                                                              
        except os.error:                                                                                                                                                                                                                                                      
            pass                                                                                                                                                                                                                                                              
        shutil.move(src, dst)                                                                                                                                                                                                                                                 
    if moves:                                                                                                                                                                                                                                                                 
        print("You might want to delete empty dirs with `find DIR -empty -type d -delete`")                                                                                                                                                                                   
                                                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                                              
def main():                                                                                                                                                                                                                                                                   
    parser = argparse.ArgumentParser(description='consolidate mergerfs')                                                                                                                                                                                                      
    parser.add_argument('disks', metavar='disk', type=str, nargs='+', help='drives to check')                                                                                                                                                                                 
    parser.add_argument('-m', dest='do_move', action='store_true', help='actual move the file')                                                                                                                                                                               
    args = parser.parse_args()                                                                                                                                                                                                                                                
    files = find(args.disks)                                                                                                                                                                                                                                                  
    moves = decide(files)                                                                                                                                                                                                                                                     
    if args.do_move:                                                                                                                                                                                                                                                          
        move(moves)                                                                                                                                                                                                                                                           
    else:                                                                                                                                                                                                                                                                     
        fake_move(moves)                                                                                                                                                                                                                                                      
                                                                                                                                                                                                                                                                              
# try with                                                                                                                                                                                                                                                                    
# $ mkdir -p mnt/disk{1,2,4}/tv/show{1,2}                                                                                                                                                                                                                                     
# $ touch mnt/disk1/tv/show1/show1.s01e01.ts mnt/disk4/tv/show1/show1.s01e02.ts mnt/disk2/tv/show2/show2.s01e01.ts mnt/disk4/tv/show2/show2.s01e02.ts                                                                                                                         
# $ python unbalancer.py mnt/disk*/tv                                                                                                                                                                                                                                         
# $ python unbalancer.py -m mnt/disk*/tv                                                                                                                                                                                                                                      
# $ find mnt -type f                                                                                                                                                                                                                                                          
# mnt/disk4/tv/show2/show2.s01e01.ts                                                                                                                                                                                                                                          
# mnt/disk4/tv/show2/show2.s01e02.ts                                                                                                                                                                                                                                          
# mnt/disk1/tv/show1/show1.s01e01.ts                                                                                                                                                                                                                                          
# mnt/disk1/tv/show1/show1.s01e02.ts                                                                                                                                                                                                                                          
                                                                                                                                                                                                                                                                              
if __name__ == "__main__":                                                                                                                                                                                                                                                    
    main()