Rewrite mkpodcastplaylist for python3 and cleanup

bin/mkpodcastplaylist

@@ -1,109 +1,119 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
-"""
-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
+from functools import partial
+from pathlib import Path
 import os
 import sys
 
+def ensure_path(p):
+    return p if isinstance(p, Path) else Path(p)
+
+def include_all(*args):
+    return True
+
 class File:
-    """A file has a size and a name"""
-
     def __init__(self, path):
-        self.name = path
+        self.path = ensure_path(path)
-        self.size = os.path.getsize(self.name)
+        self.stat = os.stat(self.path)
 
-    def get_size(self):
+    @property
-        """Get the size of the current file"""
+    def name(self):
-        return self.size
+        return self.path.name
 
-    def get_name(self):
+    def size(self):
-        """Get the name of the current file"""
+        return self.stat.st_size
-        return self.name
+
+def get(values, *, existing=set()):
+    for val in values:
+        if val in existing:
+            continue
+        if val.size() == 0:
+            continue
+        yield val
+
+def by_size(x, *, filter=None):
+    return x.size(filter=filter)
+
+def by_name(x):
+    return x.name
 
 class Directory:
-    """A directory is a list that can contain files and other directories."""
+    def __init__(self, base_dir):
+        self.base = ensure_path(base_dir)
+        self.files = {}
+        self.dirs = {}
 
-    def __init__(self, path):
+    @property
-        self.path = path
+    def name(self):
-        self.files = []
+        return self.base
 
-        files = os.listdir(path)
+    def list_by_size(self, *, files=None):
-        files.sort()
+        files = set() if files is None else files
-        for fil in files:
-            if type(fil).__name__ == "str":
-                self.files.append(File("%s/%s" % (path, fil)))
 
-    def get_dir_name(self):
+        while True:
-        """Get the name of the current directory."""
+            dirs = get(self.dirs.values(), existing=files)
-        return self.path
+            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
 
-    def get_dir_size(self):
+            files_by_name = sorted(get(self.files.values(), existing=files),
-        """Get the size of the current directory."""
+                                   key=by_name)
-        ret = 0
+            if len(files_by_name):
-        for fil in self.files:
+                current_file = files_by_name[0]
-            ret += fil.get_size()
+                files.add(current_file)
-        return ret
+                yield current_file
-
+                continue
-    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__":
+    def size(self, *, filter=None):
-    main()
+        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)