From 45c1ee5d46c6c5b4b81489e87b3fa57b6158d29d Mon Sep 17 00:00:00 2001
From: Scot Hacker <shacker@birdhouse.org>
Date: Sun, 10 Mar 2019 14:04:48 -0700
Subject: [PATCH] Handle both in-memory and local file objects

---
 todo/management/commands/import_csv.py | 11 ++--
 todo/operations/csv_importer.py        | 75 +++++++++++++-------------
 2 files changed, 47 insertions(+), 39 deletions(-)

diff --git a/todo/management/commands/import_csv.py b/todo/management/commands/import_csv.py
index e6c593d..a8b0fd1 100644
--- a/todo/management/commands/import_csv.py
+++ b/todo/management/commands/import_csv.py
@@ -1,5 +1,6 @@
 import sys
 from typing import Any
+from pathlib import Path
 
 from django.core.management.base import BaseCommand, CommandParser
 
@@ -24,11 +25,15 @@ class Command(BaseCommand):
             print("Sorry, we need a file name to work from.")
             sys.exit(1)
         else:
-            # Don't check validity of filepath here; upserter will do that.
             filepath = str(options.get("file"))
+            if not Path(filepath).exists():
+                print(f"Sorry, couldn't find file: {filepath}")
+                sys.exit(1)
 
-        importer = CSVImporter()
-        results = importer.upsert(filepath)
+        with open(filepath, mode="r", encoding="utf-8-sig") as fileobj:
+            # Pass in a file *object*, not a path
+            importer = CSVImporter()
+            results = importer.upsert(fileobj, as_string_obj=True)
 
         # Report successes, failures and summaries
         print()
diff --git a/todo/operations/csv_importer.py b/todo/operations/csv_importer.py
index afef474..a7e78a9 100644
--- a/todo/operations/csv_importer.py
+++ b/todo/operations/csv_importer.py
@@ -1,8 +1,7 @@
+import codecs
 import csv
 import datetime
 import logging
-import sys
-from pathlib import Path
 
 from django.contrib.auth import get_user_model
 from django.contrib.auth.models import Group
@@ -24,44 +23,48 @@ class CSVImporter:
         self.line_count = 0
         self.upsert_count = 0
 
-    def upsert(self, filepath):
+    def upsert(self, fileobj, as_string_obj=False):
+        """Expects a file *object*, not a file path. This is important because this has to work for both
+        the management command and the web uploader; the web uploader will pass in in-memory file
+        with no path!
 
-        if not Path(filepath).exists():
-            print(f"Sorry, couldn't find file: {filepath}")
-            sys.exit(1)
+        Header row is:
+        Title, Group, Task List, Created Date, Due Date, Completed, Created By, Assigned To, Note, Priority
+        """
 
-        with open(filepath, mode="r", encoding="utf-8-sig") as csv_file:
-            # Have arg and good file path -- read in rows as dicts.
-            # Header row is:
-            # Title, Group, Task List, Created Date, Due Date, Completed, Created By, Assigned To, Note, Priority
+        if as_string_obj:
+            # fileobj comes from mgmt command
+            csv_reader = csv.DictReader(fileobj)
+        else:
+            # fileobj comes from browser upload (in-memory)
+            csv_reader = csv.DictReader(codecs.iterdecode(fileobj, "utf-8"))
 
-            csv_reader = csv.DictReader(csv_file)
-            for row in csv_reader:
-                self.line_count += 1
+        for row in csv_reader:
+            self.line_count += 1
 
-                newrow = self.validate_row(row)
-                if newrow:
-                    # newrow at this point is fully validated, and all FK relations exist,
-                    # e.g. `newrow.get("Assigned To")`, is a Django User instance.
-                    obj, created = Task.objects.update_or_create(
-                        created_by=newrow.get("Created By"),
-                        task_list=newrow.get("Task List"),
-                        title=newrow.get("Title"),
-                        defaults={
-                            "assigned_to": newrow.get("Assigned To"),
-                            "completed": newrow.get("Completed"),
-                            "created_date": newrow.get("Created Date"),
-                            "due_date": newrow.get("Due Date"),
-                            "note": newrow.get("Note"),
-                            "priority": newrow.get("Priority"),
-                        },
-                    )
-                    self.upsert_count += 1
-                    msg = (
-                        f'Upserted task {obj.id}: "{obj.title}"'
-                        f' in list "{obj.task_list}" (group "{obj.task_list.group}")'
-                    )
-                    self.upsert_msgs.append(msg)
+            newrow = self.validate_row(row)
+            if newrow:
+                # newrow at this point is fully validated, and all FK relations exist,
+                # e.g. `newrow.get("Assigned To")`, is a Django User instance.
+                obj, created = Task.objects.update_or_create(
+                    created_by=newrow.get("Created By"),
+                    task_list=newrow.get("Task List"),
+                    title=newrow.get("Title"),
+                    defaults={
+                        "assigned_to": newrow.get("Assigned To"),
+                        "completed": newrow.get("Completed"),
+                        "created_date": newrow.get("Created Date"),
+                        "due_date": newrow.get("Due Date"),
+                        "note": newrow.get("Note"),
+                        "priority": newrow.get("Priority"),
+                    },
+                )
+                self.upsert_count += 1
+                msg = (
+                    f'Upserted task {obj.id}: "{obj.title}"'
+                    f' in list "{obj.task_list}" (group "{obj.task_list.group}")'
+                )
+                self.upsert_msgs.append(msg)
 
         self.summary_msgs.append(f"\nProcessed {self.line_count} CSV rows")
         self.summary_msgs.append(f"Upserted {self.upsert_count} rows")