Refactor for future work

1 files changed, 45 insertions, 14 deletions

diff --git a/sortashuffle.py b/sortashuffle.py
index 13a033e..31c57ef 100644
--- a/sortashuffle.py
+++ b/sortashuffle.py
@@ -21,21 +21,52 @@ TARGET  =  "<TARGET DIR>"
 SOURCES = ["<SOURCE DIR>",
            "<SOURCE DIR>"]
 
-SHOWLIST = [[f for f in os.listdir(source)] for source in SOURCES]
+def collect_showlist(sources):
+    """Returns a list of lists.
+    Each sub-list represents a show, and elements represent episodes."""
+    return [[f for f in os.listdir(source)] for source in sources]
 
-WEIGHTLIST = [len(show) for show in SHOWLIST]
+def calculate_weights(showlist):
+    """Calculate the weights for each show; # of remaining episodes"""
+    return [len(show) for show in showlist]
 
-SHUFFLED = []
-while any(SHOWLIST):
-    selection = random.choices(range(len(SHOWLIST)), weights=WEIGHTLIST, k=1)[0]
-    SHUFFLED.append(SOURCES[selection] + "//" + SHOWLIST[selection].pop(0))
-    WEIGHTLIST[selection] = len(SHOWLIST[selection])
+def select_show(showlist, weightlist):
+    """Select a show, accounting weight."""
+    return random.choices(range(len(showlist)),
+                          weights=weightlist,
+                          k=1)[0]
 
-count = 0
-for episode in SHUFFLED:
-    os.symlink(episode, TARGET + "//" + str(count))
-    count += 1
+def shuffle(showlist, weightlist):
+    """Shuffle the playlist."""
+    shuffled = []
+    while any(showlist):
+        selection = select_show(showlist, weightlist, previous)
+        shuffled.append(SOURCES[selection] + "//" + showlist[selection].pop(0))
+        weightlist[selection] = len(showlist[selection])
+    return shuffled
 
-INDEX = open(TARGET + "//" + "_PLAYLIST_INDEX.txt", 'w', encoding='utf-8')
-INDEX.write("\n".join(SHUFFLED))
-INDEX.close()
+def deploy_symlinks(shuffled):
+    """Deploy episode symlinks."""
+    count = 0
+    for episode in shuffled:
+        os.symlink(episode, TARGET + "//" + str(count))
+        count += 1
+    return count
+
+def deploy_index(shuffled):
+    """Deploy playlist index."""
+    index = open(TARGET + "//" + "_PLAYLIST_INDEX.txt", 'w', encoding='utf-8')
+    index.write("\n".join(shuffled))
+    index.close()
+    return INDEX
+
+def main():
+    SHOWLIST = collect_showlist(SOURCES)
+    WEIGHTLIST = calculate_weights(SHOWLIST)
+    SHUFFLED = shuffle(SHOWLIST, WEIGHTLIST)
+    deploy_symlinks(SHUFFLED)
+    deploy_index(SHUFFLED)
+    return 0
+
+if __name__=="__main__":
+    main()