Add option to use colors

pirate-get.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python
+from __future__ import print_function
+import __builtin__
 import webbrowser
 import urllib
 import requests
@@ -8,7 +10,6 @@ from HTMLParser import HTMLParser
 import argparse
 from pprint import pprint
 
-
 # create a subclass and override the handler methods
 class MyHTMLParser(HTMLParser):
     title = ''
@@ -45,14 +46,7 @@ def main():
     parser.add_argument('--local', dest='database', help="An xml file containing the Pirate Bay database")
     parser.add_argument('-p', dest='pages', help="The number of pages to fetch (doesn't work with --local)", default=1)
     parser.add_argument('-0', dest='first', action='store_true', help="choose the top result", default=False)
-
+    parser.add_argument('--color', dest='color', action='store_true', help="use colored output", default=False)
-    def local(args):
-        xml_str = ''
-        with open(args.database, 'r') as f:
-            xml_str += f.read()
-        htmlparser = MyHTMLParser(args.q)
-        htmlparser.feed(xml_str)
-        return htmlparser.results
 
     #todo: redo this with html parser instead of regex
     def remote(args, mirror):
@@ -91,13 +85,51 @@ def main():
                             res_l.append(curr)
                             curr = ['', 0, 0]
         except KeyboardInterrupt :
-            print "\nCancelled."
+            print("\nCancelled.")
             exit()
 
         # return the sizes in a spearate list
         return res_l, sizes, uploaded
 
     args = parser.parse_args()
+
+    def make_print():
+        if(args.color):
+            import colorama
+            colorama.init()
+            color_dict = {"default": "",
+            "header": colorama.Back.WHITE + colorama.Fore.BLACK,
+            "zebra_0": "",
+            "zebra_1": colorama.Style.DIM,
+            "WARN": colorama.Fore.YELLOW,
+            "ERROR": colorama.Fore.RED}
+            def n_print(*args, **kwargs):
+                """Print with colors"""
+                try:
+                    c = color_dict[kwargs.pop("color")]
+                    args = (c + str(args[0]),) + args[1:] + (colorama.Style.RESET_ALL,)
+                except KeyError as e:
+                    pass
+                except IndexError as e:
+                    pass
+                return __builtin__.print(*args, **kwargs)
+        else:
+            def n_print(*args, **kwargs):
+                if("color" in kwargs):
+                    kwargs.pop('color')
+                return __builtin__.print(*args, **kwargs)
+        return n_print
+
+    print=make_print()
+
+    def local(args):
+        xml_str = ''
+        with open(args.database, 'r') as f:
+            xml_str += f.read()
+        htmlparser = MyHTMLParser(args.q)
+        htmlparser.feed(xml_str)
+        return htmlparser.results
+
     if args.database:
         mags = local(args)
     else:
@@ -106,7 +138,7 @@ def main():
             res = requests.get("http://proxybay.info/list.txt").text
             mirrors += res.split("\n")[3:]
         except:
-            print "Could not fetch additional mirrors"
+            print("Could not fetch additional mirrors", color="WARN")
         for mirror in mirrors:
             try:
                 print("Trying " + mirror)
@@ -116,33 +148,45 @@ def main():
                 mags, sizes, uploaded = remote(args, mirror)
                 break
             except Exception, e:
-                print("Could not contact " + mirror)
+                print("Could not contact " + mirror, color="WARN")
 
     if mags and len(mags) > 0:
         # enhanced print output with column titles
-        print "\n%-5s %-6s %-6s %-5s %-11s %-11s  %s" % ( "LINK", "SEED", "LEECH", "RATIO", "SIZE", "UPLOAD", "NAME")
+        print("%5s %6s %6s %-5s %-11s %-11s  %s" \
-        for m in range(len(mags)):
+            % ( "LINK", "SEED", "LEECH", "RATIO", "SIZE", "UPLOAD", "NAME"),
-            magnet = mags[m]
+            color="header")
+        cur_color = "zebra_0"
+        for i in range(len(mags)):
+            magnet = mags[i]
+            no_seeders = int(magnet[1])
+            no_leechers = int(magnet[2])
             name = re.search("dn=([^\&]*)", magnet[0])
 
             # compute the S/L ratio (Higher is better)
             try:
-                ratio = float(magnet[1])/float(magnet[2])
+                ratio = no_seeders/no_leechers
             except ZeroDivisionError:
-                ratio = 0
+                ratio = -1
 
+            # Alternate between colors
+            cur_color = "zebra_0" if (cur_color == "zebra_1") else "zebra_1"
+
+            torrent_name = urllib.unquote(name.group(1).encode('ascii')) \
+                .decode('utf-8').replace("+", " ")
             # enhanced print output with justified columns
-            print "%-5s %-6s %-6s %5.1f %-11s %-11s  %s" % (m, magnet[1], magnet[2], ratio ,sizes[m], uploaded[m],urllib.unquote(name.group(1).encode('ascii')).decode('utf-8').replace("+", " ") )
+            print ("%5d %6d %6d %5.1f %-11s %-11s  %s" % (
+                i, no_seeders, no_leechers, ratio ,sizes[i],
+                uploaded[i], torrent_name), color=cur_color)
 
         if args.first:
-            print "Choosing first result";
+            print("Choosing first result");
             choice = 0
 
         else:
             try:
                 l = raw_input("Select a link: ")
             except KeyboardInterrupt :
-                print "\nCancelled."
+                print("\nCancelled.")
                 exit()
 
             try:
@@ -153,17 +197,17 @@ def main():
         if not choice == None:
             url = mags[choice][0]
             print
-            print "url:"
+            print("url:")
-            print url
+            print(url)
             if args.transmission: 
                 os.system("""transmission-remote --add "%s" """ % (url))
                 os.system("transmission-remote -l")
             else:
                 webbrowser.open(url)
         else:
-            print "Cancelled."
+            print("Cancelled.")
     else:
-        print "no results"
+        print("no results")
 
 if __name__ == "__main__":
     main()