From 3105e952ea6683ed1d16938f8a2e15cae4dea95f Mon Sep 17 00:00:00 2001 From: boukeversteegh Date: Tue, 9 Jun 2020 21:07:20 +0200 Subject: [PATCH] Fixes issue where importing cousin where path has a package with the same name broke import --- betterproto/compile/importing.py | 11 +++++------ betterproto/tests/test_get_ref_type.py | 20 ++++++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/betterproto/compile/importing.py b/betterproto/compile/importing.py index faab18f..aaf4a16 100644 --- a/betterproto/compile/importing.py +++ b/betterproto/compile/importing.py @@ -1,3 +1,4 @@ +import os import re from typing import Dict, List, Type @@ -138,16 +139,14 @@ def import_cousin(current_package, imports, py_package, py_type): package = foo.bar.baz name = foo.example.Bar """ - shared_ancestory = [ - pair[0] for pair in zip(current_package, py_package) if pair[0] == pair[1] - ] - distance_up = len(current_package) - len(shared_ancestory) + shared_ancestry = os.path.commonprefix([current_package, py_package]) + distance_up = len(current_package) - len(shared_ancestry) string_from = f".{'.' * distance_up}" + ".".join( - py_package[len(shared_ancestory) : -1] + py_package[len(shared_ancestry) : -1] ) string_import = py_package[-1] alias = f"{'_' * distance_up}" + safe_snake_case( - ".".join(py_package[len(shared_ancestory) :]) + ".".join(py_package[len(shared_ancestry) :]) ) imports.add(f"from {string_from} import {string_import} as {alias}") return f"{alias}.{py_type}" diff --git a/betterproto/tests/test_get_ref_type.py b/betterproto/tests/test_get_ref_type.py index 16a5466..b02dfea 100644 --- a/betterproto/tests/test_get_ref_type.py +++ b/betterproto/tests/test_get_ref_type.py @@ -242,6 +242,26 @@ def test_import_cousin_package(): assert name == "_y.Message" +def test_import_cousin_package_different_name(): + imports = set() + name = get_ref_type( + package="test.package1", imports=imports, source_type="cousin.package2.Message" + ) + + assert imports == {"from ...cousin import package2 as __cousin_package2"} + assert name == "__cousin_package2.Message" + + +def test_import_cousin_package_same_name(): + imports = set() + name = get_ref_type( + package="test.package", imports=imports, source_type="cousin.package.Message" + ) + + assert imports == {"from ...cousin import package as __cousin_package"} + assert name == "__cousin_package.Message" + + def test_import_far_cousin_package(): imports = set() name = get_ref_type(package="a.x.y", imports=imports, source_type="a.b.c.Message")