completed logic for wc in python

Author: TzeMingHo

implement-shell-tools/wc/custom_wc.py

@@ -2,7 +2,9 @@
 import re
 
 def format_number(number_array):
-    return [str(number).rjust(4) for number in number_array]
+    space_array = [3, 4, 4]
+    result = map(lambda number, space: str(number).rjust(space), number_array, space_array)
+    return "".join(list(result))
 
 parser = argparse.ArgumentParser(
     prog="custom wc in python",
@@ -22,7 +24,7 @@ def format_number(number_array):
 total_words = 0
 total_characters = 0
 
-numbers_row_array = []
+total_numbers_row_array = []
 
 for file_path in files_array:
     with open(file_path) as file:
@@ -31,20 +33,34 @@ def format_number(number_array):
         words = len(re.findall(r"\S+", context))
         characters = len(context)
 
+        numbers_row_array = []
         if (args.lines):
             numbers_row_array.append(lines)
         if (args.words):
             numbers_row_array.append(words)
         if (args.characters):
             numbers_row_array.append(characters)
+
+        if (len(numbers_row_array) == 1 and len(files_array) == 1):
+            print(f'{numbers_row_array[0]} {file_path}')
 
-        if (len(numbers_row_array) > 0):
-            print(numbers_row_array)
+        elif (len(numbers_row_array) > 0):
+            print(f'{format_number(numbers_row_array)} {file_path}')
         else:
-            print(lines, words, characters, file_path)
+            print(format_number([lines, words, characters]), file_path)
+
         total_lines += lines
         total_words += words
         total_characters += characters
 
 if (len(files_array) > 1):
-    print(total_lines, total_words, total_characters, 'total')
+    if (args.lines):
+        total_numbers_row_array.append(total_lines)
+    if (args.words):
+        total_numbers_row_array.append(total_words)
+    if (args.characters):
+        total_numbers_row_array.append(total_characters)
+    if (len(total_numbers_row_array) > 0):
+        print(f'{format_number(total_numbers_row_array)} total')
+    else:
+        print(format_number([total_lines, total_words, total_characters]), 'total')