diff --git a/bin/mkpodcastplaylist b/bin/mkpodcastplaylist index 3737b04..77fba95 100755 --- a/bin/mkpodcastplaylist +++ b/bin/mkpodcastplaylist @@ -1,109 +1,119 @@ -#!/usr/bin/env python -""" -This is a program that generates playlists based on the size of the -largest file/directory. It searches directories, getting their size -and lists the first file in the largest direcotry. -""" - -import errno +#!/usr/bin/env python3 + +from functools import partial +from pathlib import Path import os import sys -class File: - """A file has a size and a name""" +def ensure_path(p): + return p if isinstance(p, Path) else Path(p) + +def include_all(*args): + return True +class File: def __init__(self, path): - self.name = path - self.size = os.path.getsize(self.name) + self.path = ensure_path(path) + self.stat = os.stat(self.path) - def get_size(self): - """Get the size of the current file""" - return self.size + @property + def name(self): + return self.path.name - def get_name(self): - """Get the name of the current file""" - return self.name + def size(self): + return self.stat.st_size -class Directory: - """A directory is a list that can contain files and other directories.""" +def get(values, *, existing=set()): + for val in values: + if val in existing: + continue + if val.size() == 0: + continue + yield val - def __init__(self, path): - self.path = path - self.files = [] - - files = os.listdir(path) - files.sort() - for fil in files: - if type(fil).__name__ == "str": - self.files.append(File("%s/%s" % (path, fil))) - - def get_dir_name(self): - """Get the name of the current directory.""" - return self.path - - def get_dir_size(self): - """Get the size of the current directory.""" - ret = 0 - for fil in self.files: - ret += fil.get_size() - return ret - - def get_head_file(self): - """Get the 'head' file in the directory.""" - return self.files[0].get_name() - - def num_files(self): - """Get the number of files in the directory.""" - return len(self.files) - - def remove_head_file(self): - """Remove the head file from being counted.""" - self.files = self.files[1:] - -def sort_dir(direct): - """Sort the directories by size.""" - return direct.get_dir_size() - -def main(): - """The main script.""" - try: - base_dir = sys.argv[1] - except IndexError: - print "Must enter path of directory to search. Example: %s ./foo" % \ - sys.argv[0] - os.abort() - - if os.path.exists(base_dir) == False: - print "Could not find path: '%s'" % base_dir - os.abort() - - if base_dir[-1] == "/": - base_dir = base_dir[:-1] - - base_dir_listing = os.listdir(base_dir) - directories = [] - - for direct in base_dir_listing: - rel_dir = "%s/%s" % (base_dir, direct) - - if os.path.islink(rel_dir) == False and os.path.isdir(rel_dir): - directories.append(Directory(rel_dir)) - - while directories: - directories.sort(key=sort_dir) - directories.reverse() - headfile = directories[0].get_head_file() - #headfile = headfile.replace('podcasts/', 'podcast/', 1) - #print headfile - try: - sys.stdout.write("%s\n" % headfile) - except IOError as e: - if e.errno == errno.EPIPE: - return - raise e - directories[0].remove_head_file() - if directories[0].num_files() == 0: - directories = directories[1:] - -if __name__ == "__main__": - main() +def by_size(x, *, filter=None): + return x.size(filter=filter) + +def by_name(x): + return x.name + +class Directory: + def __init__(self, base_dir): + self.base = ensure_path(base_dir) + self.files = {} + self.dirs = {} + + @property + def name(self): + return self.base + + def list_by_size(self, *, files=None): + files = set() if files is None else files + + while True: + dirs = get(self.dirs.values(), existing=files) + without_file = lambda x: x not in files + filtered_by_size = partial(by_size, filter=without_file) + dirs_by_size = sorted(dirs, key=filtered_by_size, reverse=True) + if len(dirs_by_size): + largest_dir = dirs_by_size[0] + current_file = next(largest_dir.list_by_size(files=files)) + files.add(current_file) + yield current_file + continue + + files_by_name = sorted(get(self.files.values(), existing=files), + key=by_name) + if len(files_by_name): + current_file = files_by_name[0] + files.add(current_file) + yield current_file + continue + return + + def size(self, *, filter=None): + filter = include_all if filter is None else filter + + total = 0 + for node in self.files.values(): + if filter(node): + total += node.size() + + for node in self.dirs.values(): + total += node.size(filter=filter) + + return total + + def directory(self, *parts): + if len(parts) == 0: + return self + + part, *rest = parts + node = self.dirs.get(part) + if node is None: + node = Directory(part) + self.dirs[part] = node + + return node.directory(*rest) + + def add_file(self, path): + node = self.directory(*path.parent.parts) + node.files[path.name] = File(path) + return node.files[path.name] + +def main(base_dir): + tree = Directory(base_dir) + for root, dirs, files in os.walk(base_dir): + for f in files: + tree.add_file(Path(root) / f) + + for f in tree.list_by_size(): + print("{}/{}".format(base_dir, f.path)) + +if __name__ == '__main__': + from argparse import ArgumentParser + parser = ArgumentParser() + parser.add_argument('-o', '--output', default='.') + parser.add_argument('-d', '--directory', default=Path('.'), type=Path) + args = parser.parse_args() + main(args.directory)