From 60a8ad034fe6f173216bfad2d1d4c197d9499543 Mon Sep 17 00:00:00 2001
From: Liubov Yaronskaya <lyaronskaya@sourcerer.io>
Date: Fri, 17 Nov 2017 17:37:28 +0300
Subject: [PATCH] wip: multiple imports for opencv in python

---
 src/main/kotlin/app/extractors/PythonExtractor.kt      |  9 ++++++++-
 src/main/resources/data/imports/python.txt             |  2 ++
 src/test/kotlin/test/tests/extractors/ExtractorTest.kt | 10 ++++++++++
 3 files changed, 20 insertions(+), 1 deletion(-)
 create mode 100644 src/main/resources/data/imports/python.txt

diff --git a/src/main/kotlin/app/extractors/PythonExtractor.kt b/src/main/kotlin/app/extractors/PythonExtractor.kt
index 50e57a9d..4b7fe0be 100644
--- a/src/main/kotlin/app/extractors/PythonExtractor.kt
+++ b/src/main/kotlin/app/extractors/PythonExtractor.kt
@@ -14,6 +14,8 @@ class PythonExtractor : ExtractorInterface {
         val evaluator by lazy {
             ExtractorInterface.getLibraryClassifier(LANGUAGE_NAME)
         }
+        val MULTI_IMPORT_TO_LIB =
+                ExtractorInterface.getMultipleImportsToLibraryMap(LANGUAGE_NAME)
     }
 
     override fun extract(files: List<DiffFile>): List<CommitStats> {
@@ -35,7 +37,12 @@ class PythonExtractor : ExtractorInterface {
             }
         }
 
-        return imports.toList()
+        var libraries = imports.map { MULTI_IMPORT_TO_LIB.getOrDefault(it, it) }
+            .filter { !it.endsWith("pb")}.toMutableList()
+        if (libraries.size < imports.size) {
+            libraries.add("protobuf")
+        }
+        return libraries
 
     }
 
diff --git a/src/main/resources/data/imports/python.txt b/src/main/resources/data/imports/python.txt
new file mode 100644
index 00000000..e8a45a71
--- /dev/null
+++ b/src/main/resources/data/imports/python.txt
@@ -0,0 +1,2 @@
+cv2:opencv
+cv:opencv
\ No newline at end of file
diff --git a/src/test/kotlin/test/tests/extractors/ExtractorTest.kt b/src/test/kotlin/test/tests/extractors/ExtractorTest.kt
index ea61a440..ccc80a88 100644
--- a/src/test/kotlin/test/tests/extractors/ExtractorTest.kt
+++ b/src/test/kotlin/test/tests/extractors/ExtractorTest.kt
@@ -178,4 +178,14 @@ class ExtractorTest : Spek({
             assertExtractsImport(lib, line2, CppExtractor())
         }
     }
+
+    given("import cv2 or cv") {
+        it("imports opencv") {
+            val lib = "opencv"
+            val line1 = "import cv2"
+            assertExtractsImport(lib, line1, PythonExtractor())
+            val line2 = "import cv"
+            assertExtractsImport(lib, line2, PythonExtractor())
+        }
+    }
 })