From 754090279e515144852e5c2238062857959791f2 Mon Sep 17 00:00:00 2001
From: Matjaz Gregoric <mtyaka@gmail.com>
Date: Tue, 25 Sep 2018 10:07:09 +0200
Subject: [PATCH] Address review comments.

---
 .../commands/export_content_library.py        | 27 ++++++++++---------
 .../commands/import_content_library.py        | 17 +++++++-----
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/cms/djangoapps/contentstore/management/commands/export_content_library.py b/cms/djangoapps/contentstore/management/commands/export_content_library.py
index 76011a45cd..9c914485a1 100644
--- a/cms/djangoapps/contentstore/management/commands/export_content_library.py
+++ b/cms/djangoapps/contentstore/management/commands/export_content_library.py
@@ -3,6 +3,7 @@ Script for exporting a content library from Mongo to a tar.gz file
 """
 from __future__ import print_function
 import os
+import shutil
 
 from django.core.management.base import BaseCommand, CommandError
 from opaque_keys import InvalidKeyError
@@ -31,10 +32,9 @@ class Command(BaseCommand):
         module_store = modulestore()
         try:
             library_key = CourseKey.from_string(options['library_id'])
-            assert isinstance(library_key, LibraryLocator)
         except InvalidKeyError:
             raise CommandError(u'Invalid library ID: "{0}".'.format(options['library_id']))
-        except AssertionError:
+        if not isinstance(library_key, LibraryLocator):
             raise CommandError(u'Argument "{0}" is not a library key'.format(options['library_id']))
 
         library = module_store.get_library(library_key)
@@ -51,14 +51,15 @@ class Command(BaseCommand):
         except Exception as e:
             raise CommandError(u'Failed to export "{0}" with "{1}"'.format(library_key, e))
         else:
-            # Save generated archive with keyed filename
-            prefix, suffix, n = str(library_key).replace(':', '+'), '.tar.gz', 0
-            while os.path.exists(prefix + suffix):
-                n += 1
-                prefix = u'{0}_{1}'.format(prefix.rsplit('_', 1)[0], n) if n > 1 else u'{}_1'.format(prefix)
-            filename = prefix + suffix
-            target = os.path.join(dest_path, filename)
-            tarball.file.seek(0)
-            with open(target, 'w') as f:
-                f.write(tarball.file.read())
-            print(u'Library "{0}" exported to "{1}"'.format(library.location.library_key, target))
+            with tarball:
+                # Save generated archive with keyed filename
+                prefix, suffix, n = str(library_key).replace(':', '+'), '.tar.gz', 0
+                while os.path.exists(prefix + suffix):
+                    n += 1
+                    prefix = u'{0}_{1}'.format(prefix.rsplit('_', 1)[0], n) if n > 1 else u'{}_1'.format(prefix)
+                filename = prefix + suffix
+                target = os.path.join(dest_path, filename)
+                tarball.file.seek(0)
+                with open(target, 'w') as f:
+                    shutil.copyfileobj(tarball.file, f)
+                print(u'Library "{0}" exported to "{1}"'.format(library.location.library_key, target))
diff --git a/cms/djangoapps/contentstore/management/commands/import_content_library.py b/cms/djangoapps/contentstore/management/commands/import_content_library.py
index b967932161..014fa37946 100644
--- a/cms/djangoapps/contentstore/management/commands/import_content_library.py
+++ b/cms/djangoapps/contentstore/management/commands/import_content_library.py
@@ -12,7 +12,7 @@ from django.contrib.auth.models import User
 from django.core.exceptions import SuspiciousOperation
 from django.core.management.base import BaseCommand, CommandError
 from lxml import etree
-from opaque_keys.edx.keys import CourseKey
+from opaque_keys.edx.locator import LibraryLocator
 from path import Path
 from xmodule.contentstore.django import contentstore
 from xmodule.modulestore import ModuleStoreEnum
@@ -61,7 +61,9 @@ class Command(BaseCommand):
 
         # Gather library metadata from XML file
         xml_root = etree.parse(abs_xml_path / 'library.xml').getroot()
-        assert xml_root.tag == 'library'
+        if xml_root.tag != 'library':
+            raise CommandError(u'Failed to import {0}: Not a library archive'.format(archive_path))
+
         metadata = xml_root.attrib
         org = metadata['org']
         library = metadata['library']
@@ -73,8 +75,9 @@ class Command(BaseCommand):
 
         # Check if data would be overwritten
         ans = ''
-        while not created and ans.lower() not in ['y', 'yes', 'n', 'no']:
-            ans = raw_input(u'Library "{0}" already exists, overwrite it? [y/n] '.format(courselike_key))
+        while not created and ans not in ['y', 'yes', 'n', 'no']:
+            inp = raw_input(u'Library "{0}" already exists, overwrite it? [y/n] '.format(courselike_key))
+            ans = inp.lower()
         if ans.startswith('n'):
             print(u'Aborting import of "{0}"'.format(courselike_key))
             return
@@ -88,9 +91,9 @@ class Command(BaseCommand):
                 static_content_store=contentstore(),
                 target_id=courselike_key
             )
-        except Exception as e:
+        except Exception:
             print(u'\n=== Failed to import library-v1:{0}+{1}'.format(org, library))
-            raise e
+            raise
 
         print(u'Library "{0}" imported to "{1}"'.format(archive_path, courselike_key))
 
@@ -116,4 +119,4 @@ def _get_or_create_library(org, number, display_name, user):
         return library.location.library_key, True
     except DuplicateCourseError:
         # Course exists, return its key
-        return CourseKey.from_string(u'library-v1:{0}+{1}'.format(org, number)), False
+        return LibraryLocator(org=org, library=number), False