@@ -2,28 +2,25 @@
 #
 # Exercise 1.27
 
-import csv, sys
+import csv
+from report import read_portfolio
 
 def portfolio_cost(filename):
     'Get total portfolio cost of a csv'
 
-    total = 0
-
-    with open(filename, 'rt') as f:
-        rows = csv.reader(f)
-        next(rows)
-        for row in rows:
-            try:
-                total += int(row[1]) * float(row[2])
-            except ValueError:
-                print('Something is not valid in row: ', row)
+    total = sum([s['price']*s['shares'] for s in read_portfolio(filename)])
 
     return total
 
-if len(sys.argv) == 2:
-    filename = sys.argv[1]
-else:
-    filename = 'Data/portfolio.csv'
+def main(argv):
+    if len(argv) == 2:
+        filename = argv[1]
+    else:
+        filename = 'Data/portfolio.csv'
+
+    print('Total cost', portfolio_cost(filename))
+
 
-cost = portfolio_cost(filename)
-print('Total cost', cost)
+if __name__ == '__main__':
+    import sys
+    main(sys.argv)
@@ -2,7 +2,7 @@
 #
 # Exercise 2.4
 
-import csv, sys, locale
+import csv, locale
 from fileparse import parse_csv
 
 locale.setlocale(locale.LC_ALL, '')
@@ -75,8 +75,13 @@ def portfolio_report(portfolio_filename="Data/portfolio.csv", prices_filename="D
     print(f'{portfolio_filename:*^43s}')
     print_report(report)
 
-if len(sys.argv) > 1:
-    for file in sys.argv[1:]:
-        portfolio_report(file)
-else:
-    portfolio_report()
+def main(argv):
+    if len(argv) > 1:
+        for file in argv[1:]:
+            portfolio_report(file)
+    else:
+        portfolio_report()
+
+if __name__ == '__main__':
+    import sys
+    main(sys.argv)