Migrate import disambiguation to Python MVVM

src/main/java/com/google/api/codegen/util/py/PythonTypeTable.java

@@ -14,31 +14,43 @@
  */
 package com.google.api.codegen.util.py;
 
-import com.google.api.codegen.util.DynamicLangTypeTable;
 import com.google.api.codegen.util.NamePath;
 import com.google.api.codegen.util.TypeAlias;
 import com.google.api.codegen.util.TypeName;
 import com.google.api.codegen.util.TypeTable;
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
+import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Multimap;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.TreeMap;
 
 /** The TypeTable for Python. */
 public class PythonTypeTable implements TypeTable {
 
-  private final DynamicLangTypeTable dynamicTypeTable;
+  /** A bi-map from short names to file names. Should be kept 1:1 with moduleImports keys. */
+  private BiMap<String, String> usedShortNames = HashBiMap.create();
+
+  /**
+   * A multimap from module names to type aliases in the module. Keys should be kept 1:1 with
+   * usedShortNames.
+   */
+  private Multimap<String, TypeAlias> moduleImports = HashMultimap.create();
+
+  private final String implicitPackageName;
 
   public PythonTypeTable(String implicitPackageName) {
-    dynamicTypeTable = new DynamicLangTypeTable(implicitPackageName, ".");
+    this.implicitPackageName = implicitPackageName;
   }
 
   @Override
   public TypeTable cloneEmpty() {
-    return new PythonTypeTable(dynamicTypeTable.getImplicitPackageName());
+    return new PythonTypeTable(implicitPackageName);
   }
 
   @Override
@@ -62,7 +74,7 @@ public TypeName getTypeName(String fullName) {
 
   @Override
   public TypeName getTypeNameInImplicitPackage(String shortName) {
-    return dynamicTypeTable.getTypeNameInImplicitPackage(shortName);
+    return new TypeName(implicitPackageName + "." + shortName, shortName);
   }
 
   @Override
@@ -82,19 +94,101 @@ public String getAndSaveNicknameFor(String fullName) {
 
   @Override
   public String getAndSaveNicknameFor(TypeName typeName) {
-    return dynamicTypeTable.getAndSaveNicknameFor(typeName);
+    return typeName.getAndSaveNicknameIn(this);
   }
 
   @Override
   public String getAndSaveNicknameFor(TypeAlias alias) {
-    return dynamicTypeTable.getAndSaveNicknameFor(alias);
+    if (!alias.needsImport()) {
+      return alias.getNickname();
+    }
+
+    String shortName = alias.getNickname().substring(0, alias.getNickname().indexOf("."));
+    String className = alias.getNickname().substring(alias.getNickname().indexOf(".") + 1);
+    String moduleName =
+        alias.getFullName().substring(0, alias.getFullName().length() - className.length() - 1);
+    if (usedShortNames.containsKey(shortName)) {
+      // Short name already there.
+      String oldModuleName = usedShortNames.get(shortName);
+      if (moduleName.equals(oldModuleName)) {
+        // New alias for existing module import, no clash
+        moduleImports.put(moduleName, alias);
+        return alias.getNickname();
+      }
+
+      // Short name clashes, disambiguate.
+      String disambiguatedOldShortName = disambiguate(oldModuleName, shortName);
+      String disambiguatedNewShortName = disambiguate(moduleName, shortName);
+
+      if (!disambiguatedOldShortName.equals(disambiguatedNewShortName)) {
+        // Names were not mangled, replace old imports with disambiguated
+        usedShortNames.remove(shortName);
+        updateOldImports(disambiguatedOldShortName, moduleImports.removeAll(oldModuleName));
+      }
+
+      return getAndSaveNicknameFor(
+          TypeAlias.createAliasedImport(
+              alias.getFullName(), disambiguatedNewShortName + "." + className));
+    }
+
+    if (moduleImports.containsKey(moduleName)) {
+      // Use existing local name for already used module
+      String newShortName = usedShortNames.inverse().get(moduleName);
+      return getAndSaveNicknameFor(
+          TypeAlias.createAliasedImport(alias.getFullName(), newShortName + "." + className));
+    }
+
+    usedShortNames.put(shortName, moduleName);
+    moduleImports.put(moduleName, alias);
+    return alias.getNickname();
+  }
+
+  /* Attempts to disambiguate an import by changing the shortName by applying a number of
+   * strategies in sequence. If a strategy succeeds in modifying the shortName corresponding to the
+   * import, subsequent strategies are not attempted. In the order that they are attempted,
+   * these strategies are:
+   *
+   * Move the highest-level single package name not already present in the alias into the alias:
+   *   "from foo import bar as baz" ====> "from foo import bar as foo_baz"
+   *   "from foo.bar import baz as bar_baz" ====> "from foo.bar import baz as foo_bar_baz"
+   *
+   * Mangle with a single underscore:
+   *   "import foo as bar" ====> "import foo as bar_"
+   *   "from foo import bar as foo_bar" ====> "from foo import bar as foo_bar_"
+   *   "import foo as bar_" ====> "import foo as bar__"
+   */
+  private String disambiguate(String moduleName, String localName) {
+    List<String> nameParts = Splitter.on(".").splitToList(moduleName);
+    String localNamePackagePrefix = "";
+
+    // Move the highest-level single package name not already present in the alias into the alias.
+    for (int i = nameParts.size() - 2; i >= 0; --i) {
+      localNamePackagePrefix = nameParts.get(i) + "_" + localNamePackagePrefix;
+      if (!localName.contains(localNamePackagePrefix)) {
+        return nameParts.get(i) + "_" + localName;
+      }
+    }
+
+    // If all packages are present in the alias, mangle.
+    return localName + "_";
+  }
+
+  private void updateOldImports(String shortName, Collection<TypeAlias> aliases) {
+    for (TypeAlias alias : aliases) {
+      String className = alias.getNickname().substring(alias.getNickname().indexOf(".") + 1);
+      getAndSaveNicknameFor(
+          TypeAlias.createAliasedImport(alias.getFullName(), shortName + "." + className));
+    }
   }
 
   @Override
   public Map<String, TypeAlias> getImports() {
-    TreeMap<TypeAlias, String> inverseMap = new TreeMap<>(TypeAlias.getNicknameComparator());
-    inverseMap.putAll(dynamicTypeTable.getImportsBimap().inverse());
-    return HashBiMap.create(inverseMap).inverse();
+    ImmutableMap.Builder<String, TypeAlias> imports = ImmutableMap.builder();
+    for (Collection<TypeAlias> aliases : moduleImports.asMap().values()) {
+      TypeAlias alias = aliases.iterator().next();
+      imports.put(alias.getFullName(), alias);
+    }
+    return imports.build();
   }
 
   public boolean hasImports() {
@@ -104,6 +198,6 @@ public boolean hasImports() {
   @Override
   public String getAndSaveNicknameForInnerType(
       String containerFullName, String innerTypeShortName) {
-    return dynamicTypeTable.getAndSaveNicknameForInnerType(containerFullName, innerTypeShortName);
+    throw new UnsupportedOperationException("getAndSaveNicknameForInnerType not supported");
   }
 }

src/test/java/com/google/api/codegen/util/py/PythonTypeTableTest.java

@@ -0,0 +1,46 @@
+/* Copyright 2017 Google Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.api.codegen.util.py;
+
+import com.google.api.codegen.util.TypeAlias;
+import com.google.common.truth.Truth;
+import org.junit.Test;
+
+public class PythonTypeTableTest {
+
+  @Test
+  public void testGetAndSaveNicknameFor_disambiguate_movePackage() {
+    PythonTypeTable typeTable = new PythonTypeTable("foo.bar");
+    typeTable.getAndSaveNicknameFor(TypeAlias.create("a.c.D", "c.D"));
+    Truth.assertThat(typeTable.getAndSaveNicknameFor(TypeAlias.create("b.c.E", "c.E")))
+        .isEqualTo("b_c.E");
+  }
+
+  @Test
+  public void testGetAndSaveNicknameFor_disambiguate_move2Packages() {
+    PythonTypeTable typeTable = new PythonTypeTable("foo.bar");
+    typeTable.getAndSaveNicknameFor(TypeAlias.create("a.c.d.E", "c_d.E"));
+    Truth.assertThat(typeTable.getAndSaveNicknameFor(TypeAlias.create("b.c.d.F", "c_d.F")))
+        .isEqualTo("b_c_d.F");
+  }
+
+  @Test
+  public void testGetAndSaveNicknameFor_disambiguate_move3Packages() {
+    PythonTypeTable typeTable = new PythonTypeTable("foo.bar");
+    typeTable.getAndSaveNicknameFor(TypeAlias.create("a.c.d.e.F", "c_d_e.F"));
+    Truth.assertThat(typeTable.getAndSaveNicknameFor(TypeAlias.create("b.c.d.e.G", "c_d_e.G")))
+        .isEqualTo("b_c_d_e.G");
+  }
+}