Use contextlib to safely delete dirs in importtasks

app/tasks/importtasks.py

@@ -15,7 +15,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 
-import flask, json, os, git, tempfile, shutil, gitdb
+import flask, json, os, git, tempfile, shutil, gitdb, contextlib
 from git import GitCommandError
 from git_archive_all import GitArchiver
 from flask_sqlalchemy import SQLAlchemy
@@ -30,68 +30,7 @@ from app.tasks import celery, TaskError
 from app.utils import randomString, getExtension
 from .minetestcheck import build_tree, MinetestCheckError, ContentType
 from .minetestcheck.config import parse_conf
-
+from .krocklist import getKrockList, findModInfo
-
-krock_list_cache = None
-krock_list_cache_by_name = None
-def getKrockList():
-	global krock_list_cache
-	global krock_list_cache_by_name
-
-	if krock_list_cache is None:
-		contents = urllib.request.urlopen("https://krock-works.uk.to/minetest/modList.php").read().decode("utf-8")
-		list = json.loads(contents)
-
-		def h(x):
-			if not ("title"   in x and "author" in x and \
-					"topicId" in x and "link"   in x and x["link"] != ""):
-				return False
-
-			import re
-			m = re.search("\[([A-Za-z0-9_]+)\]", x["title"])
-			if m is None:
-				return False
-
-			x["name"] = m.group(1)
-			return True
-
-		def g(x):
-			return {
-				"title":   x["title"],
-				"author":  x["author"],
-				"name":	x["name"],
-				"topicId": x["topicId"],
-				"link":	x["link"],
-			}
-
-		krock_list_cache = [g(x) for x in list if h(x)]
-		krock_list_cache_by_name = {}
-		for x in krock_list_cache:
-			if not x["name"] in krock_list_cache_by_name:
-				krock_list_cache_by_name[x["name"]] = []
-
-			krock_list_cache_by_name[x["name"]].append(x)
-
-	return krock_list_cache, krock_list_cache_by_name
-
-
-def findModInfo(author, name, link):
-	list, lookup = getKrockList()
-
-	if name is not None and name in lookup:
-		if len(lookup[name]) == 1:
-			return lookup[name][0]
-
-		for x in lookup[name]:
-			if x["author"] == author:
-				return x
-
-	if link is not None and len(link) > 15:
-		for x in list:
-			if link in x["link"]:
-				return x
-
-	return None
 
 
 def generateGitURL(urlstr):
@@ -100,16 +39,20 @@ def generateGitURL(urlstr):
 	return "http://:@" + netloc + path + query
 
 
+@contextlib.contextmanager
 def getTempDir():
-	return os.path.join(tempfile.gettempdir(), randomString(10))
+	temp = os.path.join(tempfile.gettempdir(), randomString(10))
+	yield temp
+	shutil.rmtree(temp)
 
 
 # Clones a repo from an unvalidated URL.
 # Returns a tuple of path and repo on sucess.
 # Throws `TaskError` on failure.
 # Caller is responsible for deleting returned directory.
+@contextlib.contextmanager
 def cloneRepo(urlstr, ref=None, recursive=False):
-	gitDir = getTempDir()
+	gitDir = os.path.join(tempfile.gettempdir(), randomString(10))
 
 	err = None
 	try:
@@ -129,7 +72,9 @@ def cloneRepo(urlstr, ref=None, recursive=False):
 			for submodule in repo.submodules:
 				submodule.update(init=True)
 
-		return gitDir, repo
+		yield repo
+		shutil.rmtree(gitDir)
+		return
 
 	except GitCommandError as e:
 		# This is needed to stop the backtrace being weird
@@ -145,15 +90,12 @@ def cloneRepo(urlstr, ref=None, recursive=False):
 
 @celery.task()
 def getMeta(urlstr, author):
-	gitDir, _ = cloneRepo(urlstr, recursive=True)
+	with cloneRepo(urlstr, recursive=True) as repo:
-
 		try:
-		tree = build_tree(gitDir, author=author, repo=urlstr)
+			tree = build_tree(repo.working_tree_dir, author=author, repo=urlstr)
 		except MinetestCheckError as err:
 			raise TaskError(str(err))
 
-	shutil.rmtree(gitDir)
-
 		result = {}
 		result["name"] = tree.name
 		result["provides"] = tree.getModNames()
@@ -249,17 +191,13 @@ def updateMetaFromRelease(self, id, path):
 	print("updateMetaFromRelease: {} for {}/{}" \
 		.format(id, release.package.author.display_name, release.package.name))
 
-	temp = getTempDir()
+	with getTempDir() as temp:
-	try:
 		with ZipFile(path, 'r') as zip_ref:
 			zip_ref.extractall(temp)
 
 		postReleaseCheckUpdate(self, release, temp)
 		db.session.commit()
 
-	finally:
-		shutil.rmtree(temp)
-
 
 @celery.task(bind=True)
 def checkZipRelease(self, id, path):
@@ -269,8 +207,7 @@ def checkZipRelease(self, id, path):
 	elif release.package is None:
 		raise TaskError("No package attached to release")
 
-	temp = getTempDir()
+	with getTempDir() as temp:
-	try:
 		with ZipFile(path, 'r') as zip_ref:
 			zip_ref.extractall(temp)
 
@@ -280,9 +217,6 @@ def checkZipRelease(self, id, path):
 		release.approve(release.package.author)
 		db.session.commit()
 
-	finally:
-		shutil.rmtree(temp)
-
 
 @celery.task(bind=True)
 def makeVCSRelease(self, id, branch):
@@ -292,16 +226,14 @@ def makeVCSRelease(self, id, branch):
 	elif release.package is None:
 		raise TaskError("No package attached to release")
 
-	gitDir, repo = cloneRepo(release.package.repo, ref=branch, recursive=True)
+	with cloneRepo(release.package.repo, ref=branch, recursive=True) as repo:
+		postReleaseCheckUpdate(self, release, repo.working_tree_dir)
 
-	postReleaseCheckUpdate(self, release, gitDir)
-
-	try:
 		filename = randomString(10) + ".zip"
 		destPath = os.path.join(app.config["UPLOAD_DIR"], filename)
 
 		assert(not os.path.isfile(destPath))
-		archiver = GitArchiver(force_sub=True, main_repo_abspath=gitDir)
+		archiver = GitArchiver(force_sub=True, main_repo_abspath=repo.working_tree_dir)
 		archiver.create(destPath)
 		assert(os.path.isfile(destPath))
 
@@ -314,8 +246,6 @@ def makeVCSRelease(self, id, branch):
 		updateMetaFromRelease.delay(release.id, destPath)
 
 		return release.url
-	finally:
-		shutil.rmtree(gitDir)
 
 
 @celery.task()
@@ -324,18 +254,10 @@ def importRepoScreenshot(id):
 	if package is None or package.state == PackageState.DELETED:
 		raise Exception("Unexpected none package")
 
-	# Get URL Maker
-	try:
-		gitDir, _ = cloneRepo(package.repo)
-	except TaskError as e:
-		# ignore download errors
-		print(e)
-		return None
-
-	# Find and import screenshot
 	try:
+		with cloneRepo(package.repo) as repo:
 			for ext in ["png", "jpg", "jpeg"]:
-			sourcePath = gitDir + "/screenshot." + ext
+				sourcePath = repo.working_tree_dir + "/screenshot." + ext
 				if os.path.isfile(sourcePath):
 					filename = randomString(10) + "." + ext
 					destPath = os.path.join(app.config["UPLOAD_DIR"], filename)
@@ -350,8 +272,11 @@ def importRepoScreenshot(id):
 					db.session.commit()
 
 					return "/uploads/" + filename
-	finally:
+
-		shutil.rmtree(gitDir)
+	except TaskError as e:
+		# ignore download errors
+		print(e)
+		pass
 
 	print("screenshot.png does not exist")
 	return None

app/tasks/krocklist.py

@@ -0,0 +1,80 @@
+# ContentDB
+# Copyright (C) 2018-20 rubenwardy
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+
+import json
+import urllib.request
+
+krock_list_cache = None
+krock_list_cache_by_name = None
+def getKrockList():
+	global krock_list_cache
+	global krock_list_cache_by_name
+
+	if krock_list_cache is None:
+		contents = urllib.request.urlopen("https://krock-works.uk.to/minetest/modList.php").read().decode("utf-8")
+		list = json.loads(contents)
+
+		def h(x):
+			if not ("title"   in x and "author" in x and \
+					"topicId" in x and "link"   in x and x["link"] != ""):
+				return False
+
+			import re
+			m = re.search(r"\[([A-Za-z0-9_]+)\]", x["title"])
+			if m is None:
+				return False
+
+			x["name"] = m.group(1)
+			return True
+
+		def g(x):
+			return {
+				"title":   x["title"],
+				"author":  x["author"],
+				"name":	x["name"],
+				"topicId": x["topicId"],
+				"link":	x["link"],
+			}
+
+		krock_list_cache = [g(x) for x in list if h(x)]
+		krock_list_cache_by_name = {}
+		for x in krock_list_cache:
+			if not x["name"] in krock_list_cache_by_name:
+				krock_list_cache_by_name[x["name"]] = []
+
+			krock_list_cache_by_name[x["name"]].append(x)
+
+	return krock_list_cache, krock_list_cache_by_name
+
+
+def findModInfo(author, name, link):
+	list, lookup = getKrockList()
+
+	if name is not None and name in lookup:
+		if len(lookup[name]) == 1:
+			return lookup[name][0]
+
+		for x in lookup[name]:
+			if x["author"] == author:
+				return x
+
+	if link is not None and len(link) > 15:
+		for x in list:
+			if link in x["link"]:
+				return x
+
+	return None