Merge pull request #21 from neilccbrown/import-as

Added support for "import x as y" in .pyi files

Author: Tobias-Kohn

build.sbt

@@ -1,7 +1,7 @@
 import sbtcrossproject.CrossPlugin.autoImport.{crossProject, CrossType}
 
 val useScalaVersion = "2.13.14"
-val releaseVersion = "1.1.0"
+val releaseVersion = "1.1.1"
 
 val sharedSettings = Seq(
   scalaVersion := useScalaVersion,

package-lock.json

@@ -1,7 +1,7 @@
 
 {
   "name": "tigerpython-parser",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "description": "Enhanced error recognition in Python ",
   "main": "release/tigerpython-parser.mjs",
   "types": "tpParser/js/types/index.d.ts",

package.json

@@ -3,12 +3,18 @@ package tigerpython.utilities.completer
 import tigerpython.utilities.fastparse._
 import tigerpython.utilities.types._
 
+import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
 /**
  * This takes a Pyi-file and defines a TigerPython-module from it.
  */
-class PyiModuleParser(val module: Module) extends PyiParser {
+class PyiModuleParser(val module: Module, val moduleLookup: mutable.Map[String, DataType]) extends PyiParser {
+
+  // Maps an alias to a module, imported either as:
+  // import math          # produces "math" -> <math Module>
+  // import math as m     # produces "m" -> <math Module>
+  private val fullModuleImports: mutable.Map[String, Module] = mutable.Map()
 
   private var currentClass: PythonClass = _
 
@@ -36,6 +42,17 @@ class PyiModuleParser(val module: Module) extends PyiParser {
         for (el <- elts.tail)
           tp = DataType.getCompatibleType(tp, convertToType(el))
         tp
+      case AttributeNode(base, name) =>
+        val baseDotted = PyiModuleParser.toDotted(base)
+        if (baseDotted != null && fullModuleImports.contains(baseDotted))
+          fullModuleImports(baseDotted).findField(name) match {
+            case Some (tp) =>
+              tp
+            case _ =>
+              BuiltinTypes.ANY_TYPE
+          }
+        else
+          BuiltinTypes.ANY_TYPE
       case _ =>
         BuiltinTypes.ANY_TYPE
     }
@@ -113,8 +130,56 @@ class PyiModuleParser(val module: Module) extends PyiParser {
   }
 
   override protected
-  def importModule(module: String, alias: String): Boolean = false
+  def importModule(module: String, alias: String): Boolean = {
+    val segments = module.split("\\.")
+    if (segments.isEmpty) return false
+
+    // Start with the outermost module
+    var currentOpt = moduleLookup.get(segments.head) match {
+      case Some(m: Module) if m != null => Some(m)
+      case _ => None
+    }
+
+    // Traverse nested modules
+    for (name <- segments.tail if currentOpt.isDefined) {
+      currentOpt = currentOpt match {
+        case Some(mod) =>
+          mod.getFields(name) match {
+            case m: Module if m != null => Some(m)
+            case _ => None
+          }
+        case None => None
+      }
+    }
+
+    currentOpt match {
+      case Some(finalModule) =>
+        if (alias == null)
+          fullModuleImports(module) = finalModule
+        else
+          fullModuleImports(alias) = finalModule
+        true
+      case None =>
+        false
+    }
+  }
+
 
   override protected
   def importNameFromModule(module: String, name: String, alias: String): Boolean = false
 }
+
+object PyiModuleParser {
+  private def toDotted(exprAst: ExprAst): String = {
+    exprAst match {
+      case NameNode(name) =>
+        name
+      case AttributeNode(base, name) =>
+        val prefix = toDotted(base)
+        if (prefix == null) null
+        else s"$prefix.$name"
+      case _ =>
+        null
+    }
+  }
+}

release/tigerpython-parser.js

@@ -227,7 +227,7 @@ object ModuleLoader {
       val names = name.split('.')
       val lastName = names(names.length - 1)
       val module = new Module(lastName)
-      val parser = new PyiModuleParser(module)
+      val parser = new PyiModuleParser(module, modules)
       parser.parse(body)
       if (names.length == 1) {
         modules(name) = module

release/tigerpython-parser.mjs

@@ -9,6 +9,11 @@ package tigerpython.utilities.types
 object BuiltinTypes {
   val builtins: collection.mutable.Map[String, DataType] = collection.mutable.Map[String, DataType]()
 
+  def add(types: Instance*): Unit =
+    for (t <- types)
+      builtins(t.baseType.name) = t.baseType
+
+
   def fromString(s: String): DataType = {
     val result = builtins.getOrElse(s, ExceptionTypes.byName(s))
     if (result != null)
@@ -294,6 +299,8 @@ object BuiltinTypes {
   val STRING = new Instance(STRING_TYPE)
   val TUPLE = new Instance(TUPLE_TYPE)
 
+  add(BOOLEAN, BUFFER, BYTEARRAY, COMPLEX, DICT, FILE, FLOAT, FROZENSET, GENERATOR, INTEGER, ITERATOR, LIST, LONG, NONE, SET, STRING, TUPLE);
+
   def appyTigerPythonModifications(): Unit = {
     LIST_TYPE.addFields(
       PrimitiveType("first"),

tpParser/shared/src/main/scala/tigerpython/utilities/completer/PyiModuleParser.scala

@@ -0,0 +1,3 @@
+# 6
+# capitalize;center;count;decode;encode;endswith;expandtabs;find;format;index;isalnum;isalpha;isdigit;islower;isspace;istitle;isupper;join;ljust;lower;lstrip;partition;replace;rfind;rindex;rjust;rpartition;rsplit;rstrip;split;splitlines;startswith;strip;swapcase;title;translate;upper;zfill
+"foo".

tpParser/shared/src/main/scala/tigerpython/utilities/scopes/ModuleLoader.scala

@@ -0,0 +1,3 @@
+# 14
+# capitalize;center;count;decode;encode;endswith;expandtabs;find;format;index;isalnum;isalpha;isdigit;islower;isspace;istitle;isupper;join;ljust;lower;lstrip;partition;replace;rfind;rindex;rjust;rpartition;rsplit;rstrip;split;splitlines;startswith;strip;swapcase;title;translate;upper;zfill
+"foo".upper().