Merge pull request #205 from Feuermurmel/list-fix

Refactoring in the lists module and fix of #198

src/lists.py

@@ -35,61 +35,75 @@ File listing functions.
 
 
 from ucs import *
+import itertools
 
 
-def _columnise(ponies):
+def _columnise_list(items: list, available_width: int, length_fn: callable, separation : int = 2):
+    """
+    From a list of items, produce a list of columns. Each columns is a list of tuples. Each tuple contains the element and a an int, specifying how much shorter the element is compared to the column width. 
+    
+    :param items: A list of items.
+    :param available_width: The maximum with available horizontally for the items.
+    :param length_fn: Function that is called to compute an element's width.
+    :param separation: Amount of space to insert between columns. 
+    """
+    
+    num_items = len(items)
+    items_with_length = [(i, length_fn(i)) for i in items]
+    max_item_length = max(i for _, i in items_with_length)
+    
+    # Make at least one column to handle very narrow terminals.
+    num_columns = max((available_width + separation) // (max_item_length + separation), 1)
+    column_length = (num_items - 1) // num_columns + 1
+    
+    items_with_spacing = [(i, max_item_length - l + separation) for i, l in items_with_length]
+    
+    return [items_with_spacing[i:i + column_length] for i in range(0, num_items, column_length)]
+
+
+def _print_columnised(items):
     '''
     Columnise a list and prints it
     
-    @param  ponies:list<(str, str)>  All items to list, each item should have to elements: unformated name, formated name
+    @param  ponies:list<(str, str)>  All items to list, each item should have to elements: unformatted name, formatted name.
     '''
-    ## Get terminal width, and a 2 which is the space between columns
+    ## Get terminal width
-    termwidth = gettermsize()[1] + 2
+    _, term_width = gettermsize()
-    ## Sort the ponies, and get the cells' widths, and the largest width + 2
-    ponies.sort(key = lambda pony : pony[0])
-    widths = [UCS.dispLen(pony[0]) for pony in ponies]
-    width = max(widths) + 2 # longest pony file name + space between columns
     
-    ## Calculate the number of rows and columns, can create a list of empty columns
+    columns = _columnise_list(
-    cols = termwidth // width # do not believe electricians, this means ⌊termwidth / width⌋
+        sorted(items, key = lambda x: x[0]),
-    rows = (len(ponies) + cols - 1) // cols
+        term_width,
-    columns = []
+        lambda x: UCS.dispLen(x[0]))
-    for c in range(0, cols):  columns.append([])
     
-    ## Fill the columns with cells of ponies
+    for row in itertools.zip_longest(*columns):
-    (y, x) = (0, 0)
+        def iter_parts():
-    for j in range(0, len(ponies)):
+            spacing = 0
-        cell = ponies[j][1] + ' ' * (width - widths[j]);
-        columns[x].append(cell)
-        y += 1
-        if y == rows:
-            x += 1
-            y = 0
             
-    ## Make the columnisation nicer by letting the last row be partially empty rather than the last column
+            for cell in row:
-    diff = rows * cols - len(ponies)
+                if cell:
-    if (diff > 2) and (rows > 1):
+                    # Yield this here to prevent whitespace to be printed after the last column.
-        c = cols - 1
+                    yield ' ' * spacing
-        diff -= 1
-        while diff > 0:
-            columns[c] = columns[c - 1][-diff:] + columns[c]
-            c -= 1
-            columns[c] = columns[c][:-diff]
-            diff -= 1
                     
-    ## Create rows from columns
+                    (_, item), spacing = cell
-    lines = []
+                    
-    for r in range(0, rows):
+                    yield item
-         lines.append([])
+        
-         for c in range(0, cols):
+        print(''.join(iter_parts()))
-             if r < len(columns[c]):
-                 line = lines[r].append(columns[c][r])
     
-    ## Print the matrix, with one extra blank row
-    print('\n'.join([''.join(line)[:-2] for line in lines]))
     print()
 
 
+def _get_file_list(dir_path : str, extension : str):
+    """
+    Return a list of files in the specified directory with have the specified file name extension.
+    
+    :param dir_path: Path to the directory.
+    :param extension: The allowed file name extension.
+    """
+    
+    return [i[:-len(extension)] for i in os.listdir(dir_path) if endswith(i, extension)]
+
+
 def simplelist(ponydirs, quoters = [], ucsiser = None):
     '''
     Lists the available ponies
@@ -100,13 +114,9 @@ def simplelist(ponydirs, quoters = [], ucsiser = None):
     '''
     for ponydir in ponydirs: # Loop ponydirs
         ## Get all ponies in the directory
-        _ponies = os.listdir(ponydir)
+        ponies = _get_file_list(ponydir, '.pony')
         
-        ## Remove .pony from all files and skip those that does not have .pony
+        print()
-        ponies = []
-        for pony in _ponies:
-            if endswith(pony, '.pony'):
-                ponies.append(pony[:-5])
         
         ## UCS:ise pony names, they are already sorted
         if ucsiser is not None:
@@ -116,7 +126,7 @@ def simplelist(ponydirs, quoters = [], ucsiser = None):
         if len(ponies) == 0:
             continue
         print('\033[1mponies located in ' + ponydir + '\033[21m')
-        _columnise([(pony, '\033[1m' + pony + '\033[21m' if pony in quoters else pony) for pony in ponies])
+        _print_columnised([(pony, '\033[1m' + pony + '\033[21m' if pony in quoters else pony) for pony in ponies])
 
 
 def linklist(ponydirs = None, quoters = [], ucsiser = None):
@@ -127,18 +137,10 @@ def linklist(ponydirs = None, quoters = [], ucsiser = None):
     @param  quoters:__in__(str)→bool                  Set of ponies that of quotes
     @param  ucsiser:(list<str>, map<str, str>)?→void  Function used to UCS:ise names
     '''
-    ## Get the size of the terminal
-    termsize = gettermsize()
     
     for ponydir in ponydirs: # Loop ponydirs
         ## Get all pony files in the directory
-        _ponies = os.listdir(ponydir)
+        ponies = _get_file_list(ponydirs, '.pony')
-        
-        ## Remove .pony from all files and skip those that does not have .pony
-        ponies = []
-        for pony in _ponies:
-            if endswith(pony, '.pony'):
-                ponies.append(pony[:-5])
         
         ## If there are no ponies in the directory skip to next directory, otherwise, print the directories name
         if len(ponies) == 0:
@@ -193,7 +195,7 @@ def linklist(ponydirs = None, quoters = [], ucsiser = None):
             ponies[(item.replace('\033[1m', '').replace('\033[21m', ''), item)] = w
         
         ## Print the ponies, columnised
-        _columnise(list(ponies))
+        _print_columnised(list(ponies))
 
 
 def onelist(standarddirs, extradirs = None, ucsiser = None):
@@ -205,19 +207,7 @@ def onelist(standarddirs, extradirs = None, ucsiser = None):
     @param  ucsiser:(list<str>)?→void  Function used to UCS:ise names
     '''
     ## Get all pony files
-    _ponies = []
+    ponies = [j for i in [standarddirs, extradirs] for j in _get_file_list(i, '.pony')]
-    if standarddirs is not None:
-        for ponydir in standarddirs:
-            _ponies += os.listdir(ponydir)
-    if extradirs is not None:
-        for ponydir in extradirs:
-            _ponies += os.listdir(ponydir)
-        
-    ## Remove .pony from all files and skip those that does not have .pony
-    ponies = []
-    for pony in _ponies:
-        if endswith(pony, '.pony'):
-            ponies.append(pony[:-5])
     
     ## UCS:ise and sort
     if ucsiser is not None:
@@ -239,24 +229,11 @@ def balloonlist(balloondirs, isthink):
     @param  balloondirs:itr<str>  The balloon directories to use
     @param  isthink:bool          Whether the ponythink command is used
     '''
-    ## Get the size of the terminal
+    
-    termsize = gettermsize()
+    extension = '.think' if isthink else '.say'
     
     ## Get all balloons
-    balloonset = set()
+    balloonset = set(j for i in balloondirs for j in _get_file_list(i, extension))
-    for balloondir in balloondirs:
-        for balloon in os.listdir(balloondir):
-            ## Use .think if running ponythink, otherwise .say
-            if isthink and endswith(balloon, '.think'):
-                balloon = balloon[:-6]
-            elif (not isthink) and endswith(balloon, '.say'):
-                balloon = balloon[:-4]
-            else:
-                continue
-            
-            ## Add the balloon if there is none with the same name
-            if balloon not in balloonset:
-                balloonset.add(balloon)
     
     ## Print all balloos, columnised
-    _columnise([(balloon, balloon) for balloon in balloonset])
+    _print_columnised([(balloon, balloon) for balloon in balloonset])