From 6adf0581dc70316610af9d2f0d4c6814e35aac50 Mon Sep 17 00:00:00 2001
From: Viktor Stanchev <me@viktorstanchev.com>
Date: Sat, 19 Sep 2015 23:29:12 -0700
Subject: [PATCH] first step towards prettytable

not pretty yet
---
 pirate/print.py     | 38 +++++++++++++++++++++-----------------
 setup.py            |  2 +-
 tests/test_print.py | 31 ++++++++++++++++++++++---------
 3 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/pirate/print.py b/pirate/print.py
index a83b89b..eac3f1b 100644
--- a/pirate/print.py
+++ b/pirate/print.py
@@ -2,12 +2,14 @@ import builtins
 import re
 import os
 import gzip
-import colorama
 import urllib.parse as parse
 import urllib.request as request
 import shutil
 from io import BytesIO
 
+import colorama
+import prettytable
+
 import pirate.data
 
 
@@ -32,20 +34,25 @@ def print(*args, **kwargs):
         return builtins.print(*args, **kwargs)
 
 
-# TODO: extract the name from the search results instead of the magnet link when possible
+# TODO: extract the name from the search results instead of from the magnet link when possible
 def search_results(results, local=None):
     columns = shutil.get_terminal_size((80, 20)).columns
     cur_color = 'zebra_0'
 
     if local:
-        print('{:>4}   {:{length}}'.format(
-              'LINK', 'NAME', length=columns - 8),
-              color='header')
+        table = prettytable.PrettyTable(['LINK', 'NAME'])
     else:
-        print('{:>4}  {:>5}  {:>5}  {:>5}  {:9}  {:11}  {:{length}}'.format(
-              'LINK', 'SEED', 'LEECH', 'RATIO',
-              'SIZE', 'UPLOAD', 'NAME', length=columns - 52),
-              color='header')
+        table = prettytable.PrettyTable(['LINK', 'SEED', 'LEECH', 'RATIO', 'SIZE', '', 'UPLOAD', 'NAME'])
+        table.align['NAME'] = 'l'
+        table.align['SEED'] = 'r'
+        table.align['LEECH'] = 'r'
+        table.align['RATIO'] = 'r'
+        table.align['SIZE'] = 'r'
+        table.align['UPLOAD'] = 'l'
+
+    table.max_width = columns
+    table.border = False
+    table.padding_width = 1
 
     for n, result in enumerate(results):
         # Alternate between colors
@@ -55,8 +62,7 @@ def search_results(results, local=None):
         torrent_name = parse.unquote_plus(name.group(1))
 
         if local:
-            line = '{:5}  {:{length}}'
-            content = [n, torrent_name[:columns]]
+            content = [n, torrent_name[:columns - 7]]
         else:
             no_seeders = int(result['seeds'])
             no_leechers = int(result['leechers'])
@@ -74,13 +80,11 @@ def search_results(results, local=None):
             except ZeroDivisionError:
                 ratio = float('inf')
 
-            line = ('{:4}  {:5}  {:5}  {:5.1f}  {:5.1f}'
-                    ' {:3}  {:<11}  {:{length}}')
-            content = [n, no_seeders, no_leechers, ratio,
-                       size, unit, date, torrent_name[:columns - 52]]
+            content = [n, no_seeders, no_leechers, '{0:.1f}'.format(ratio),
+                       '{0:.1f}'.format(size), unit, date, torrent_name[:columns - 53]]
 
-        # enhanced print output with justified columns
-        print(line.format(*content, length=columns - 52), color=cur_color)
+        table.add_row(content)
+    print(table)
 
 
 def descriptions(chosen_links, results, site):
diff --git a/setup.py b/setup.py
index 8d30a87..0af554d 100755
--- a/setup.py
+++ b/setup.py
@@ -13,7 +13,7 @@ setup(name='pirate-get',
     entry_points={
         'console_scripts': ['pirate-get = pirate.pirate:main']
     },
-    install_requires=['colorama>=0.3.3', 'pyquery>=1.2.9'],
+    install_requires=['colorama>=0.3.3', 'pyquery>=1.2.9', 'PrettyTable>=0.7.2'],
     keywords=['torrent', 'magnet', 'download', 'tpb', 'client'],
     classifiers=[
         'Topic :: Utilities',
diff --git a/tests/test_print.py b/tests/test_print.py
index 438170d..5795b5d 100755
--- a/tests/test_print.py
+++ b/tests/test_print.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 import unittest
-from unittest.mock import patch
-from unittest.mock import call
+from unittest.mock import patch, call, MagicMock
 
 import pirate.print
 
@@ -9,7 +8,11 @@ import pirate.print
 class TestPrint(unittest.TestCase):
 
     def test_print_results(self):
-        with patch('pirate.print.print') as mock:
+        class MockTable:
+            add_row = MagicMock()
+            align = {}
+        mock = MockTable()
+        with patch('prettytable.PrettyTable', return_value=mock) as prettytable:
             results = [{
                 'magnet': 'dn=name',
                 'seeds': 1,
@@ -18,12 +21,22 @@ class TestPrint(unittest.TestCase):
                 'uploaded': 'never'
             }]
             pirate.print.search_results(results)
-            actual = mock.call_args_list
-            expected = [
-                call('LINK   SEED  LEECH  RATIO  SIZE       UPLOAD       NAME                        ', color='header'),
-                call('   0      1      2    0.5    3.0 MiB  never        name                        ', color='zebra_1'),
-            ]
-            self.assertEqual(expected, actual)
+            prettytable.assert_called_once_with(['LINK', 'SEED', 'LEECH', 'RATIO', 'SIZE', '', 'UPLOAD', 'NAME'])
+            mock.add_row.assert_has_calls([call([0, 1, 2, '0.5', '3.0', 'MiB', 'never', 'name'])])
+
+    def test_print_results(self):
+        class MockTable:
+            add_row = MagicMock()
+            align = {}
+        mock = MockTable()
+        with patch('prettytable.PrettyTable', return_value=mock) as prettytable:
+            results = [{
+                'magnet': 'dn=name',
+                'Name': 'name',
+            }]
+            pirate.print.search_results(results, local=True)
+            prettytable.assert_called_once_with(['LINK', 'NAME'])
+            mock.add_row.assert_has_calls([call([0, 'name'])])
 
 if __name__ == '__main__':
     unittest.main()