Skip to content

fix: Harden resolve_import with new changes #803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
EdwardJXLi merged 1 commit into from
Mar 12, 2025
Merged

fix: Harden resolve_import with new changes #803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
194 changes: 99 additions & 95 deletions src/codegen/sdk/python/import_resolution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,73 +85,87 @@ def imported_exports(self) -> list[Exportable]:
@noapidoc
@reader
def resolve_import(self, base_path: str | None = None, *, add_module_name: str | None = None) -> ImportResolution[PyFile] | None:
base_path = base_path or self.ctx.projects[0].base_path or ""
module_source = self.module.source if self.module else ""
symbol_name = self.symbol_name.source if self.symbol_name else ""
if add_module_name:
module_source += f".{symbol_name}"
symbol_name = add_module_name
# If import is relative, convert to absolute path
if module_source.startswith("."):
module_source = self._relative_to_absolute_import(module_source)

# =====[ Check if we are importing an entire file ]=====
if self.is_module_import():
# covers `import a.b.c` case and `from a.b.c import *` case
filepath = os.path.join(base_path, module_source.replace(".", "/") + ".py")
else:
# This is the case where you do:
# `from a.b.c import foo`
filepath = os.path.join(
base_path,
module_source.replace(".", "/") + "/" + symbol_name + ".py",
)

# =====[ Check if we are importing an entire file with custom resolve path or sys.path enabled ]=====
if len(self.ctx.config.import_resolution_paths) > 0 or self.ctx.config.py_resolve_syspath:
# Handle resolve overrides first if both is set
resolve_paths: list[str] = self.ctx.config.import_resolution_paths + (sys.path if self.ctx.config.py_resolve_syspath else [])
if file := self._file_by_custom_resolve_paths(resolve_paths, filepath):
try:
base_path = base_path or self.ctx.projects[0].base_path or ""
module_source = self.module.source if self.module else ""
symbol_name = self.symbol_name.source if self.symbol_name else ""
if add_module_name:
module_source += f".{symbol_name}"
symbol_name = add_module_name
# If import is relative, convert to absolute path
if module_source.startswith("."):
module_source = self._relative_to_absolute_import(module_source)

# =====[ Check if we are importing an entire file ]=====
if self.is_module_import():
# covers `import a.b.c` case and `from a.b.c import *` case
filepath = os.path.join(base_path, module_source.replace(".", "/") + ".py")
else:
# This is the case where you do:
# `from a.b.c import foo`
filepath = os.path.join(
base_path,
module_source.replace(".", "/") + "/" + symbol_name + ".py",
)

# =====[ Check if we are importing an entire file with custom resolve path or sys.path enabled ]=====
if len(self.ctx.config.import_resolution_paths) > 0 or self.ctx.config.py_resolve_syspath:
# Handle resolve overrides first if both is set
resolve_paths: list[str] = self.ctx.config.import_resolution_paths + (sys.path if self.ctx.config.py_resolve_syspath else [])
if file := self._file_by_custom_resolve_paths(resolve_paths, filepath):
return ImportResolution(from_file=file, symbol=None, imports_file=True)

# =====[ Default path ]=====
if file := self.ctx.get_file(filepath):
return ImportResolution(from_file=file, symbol=None, imports_file=True)

filepath = filepath.replace(".py", "/__init__.py")
if file := self.ctx.get_file(filepath):
# TODO - I think this is another edge case, due to `dao/__init__.py` etc.
# You can't do `from a.b.c import foo` => `foo.utils.x` right now since `foo` is just a file...
return ImportResolution(from_file=file, symbol=None, imports_file=True)

# =====[ Default path ]=====
if file := self.ctx.get_file(filepath):
return ImportResolution(from_file=file, symbol=None, imports_file=True)

filepath = filepath.replace(".py", "/__init__.py")
if file := self.ctx.get_file(filepath):
# TODO - I think this is another edge case, due to `dao/__init__.py` etc.
# You can't do `from a.b.c import foo` => `foo.utils.x` right now since `foo` is just a file...
return ImportResolution(from_file=file, symbol=None, imports_file=True)

# =====[ Check if `module.py` file exists in the graph with custom resolve path or sys.path enabled ]=====
filepath = module_source.replace(".", "/") + ".py"
if len(self.ctx.config.import_resolution_paths) > 0 or self.ctx.config.py_resolve_syspath:
# Handle resolve overrides first if both is set
resolve_paths: list[str] = self.ctx.config.import_resolution_paths + (sys.path if self.ctx.config.py_resolve_syspath else [])
if file := self._file_by_custom_resolve_paths(resolve_paths, filepath):
# =====[ Check if `module.py` file exists in the graph with custom resolve path or sys.path enabled ]=====
filepath = module_source.replace(".", "/") + ".py"
if len(self.ctx.config.import_resolution_paths) > 0 or self.ctx.config.py_resolve_syspath:
# Handle resolve overrides first if both is set
resolve_paths: list[str] = self.ctx.config.import_resolution_paths + (sys.path if self.ctx.config.py_resolve_syspath else [])
if file := self._file_by_custom_resolve_paths(resolve_paths, filepath):
symbol = file.get_node_by_name(symbol_name)
return ImportResolution(from_file=file, symbol=symbol)

# =====[ Check if `module.py` file exists in the graph ]=====
filepath = os.path.join(base_path, filepath)
if file := self.ctx.get_file(filepath):
symbol = file.get_node_by_name(symbol_name)
return ImportResolution(from_file=file, symbol=symbol)

# =====[ Check if `module.py` file exists in the graph ]=====
filepath = os.path.join(base_path, filepath)
if file := self.ctx.get_file(filepath):
symbol = file.get_node_by_name(symbol_name)
if symbol is None:
if file.get_node_from_wildcard_chain(symbol_name):
return ImportResolution(from_file=file, symbol=None, imports_file=True)
if symbol is None:
if file.get_node_from_wildcard_chain(symbol_name):
return ImportResolution(from_file=file, symbol=None, imports_file=True)
else:
# This is most likely a broken import
return ImportResolution(from_file=file, symbol=None)
else:
# This is most likely a broken import
return ImportResolution(from_file=file, symbol=None)
else:
return ImportResolution(from_file=file, symbol=symbol)

# =====[ Check if `module/__init__.py` file exists in the graph with custom resolve path or sys.path enabled ]=====
filepath = filepath.replace(".py", "/__init__.py")
if len(self.ctx.config.import_resolution_paths) > 0 or self.ctx.config.py_resolve_syspath:
# Handle resolve overrides first if both is set
resolve_paths: list[str] = self.ctx.config.import_resolution_paths + (sys.path if self.ctx.config.py_resolve_syspath else [])
if from_file := self._file_by_custom_resolve_paths(resolve_paths, filepath):
return ImportResolution(from_file=file, symbol=symbol)

# =====[ Check if `module/__init__.py` file exists in the graph with custom resolve path or sys.path enabled ]=====
filepath = filepath.replace(".py", "/__init__.py")
if len(self.ctx.config.import_resolution_paths) > 0 or self.ctx.config.py_resolve_syspath:
# Handle resolve overrides first if both is set
resolve_paths: list[str] = self.ctx.config.import_resolution_paths + (sys.path if self.ctx.config.py_resolve_syspath else [])
if from_file := self._file_by_custom_resolve_paths(resolve_paths, filepath):
symbol = from_file.get_node_by_name(symbol_name)
if symbol is None:
if from_file.get_node_from_wildcard_chain(symbol_name):
return ImportResolution(from_file=from_file, symbol=None, imports_file=True)
else:
# This is most likely a broken import
return ImportResolution(from_file=from_file, symbol=None)

else:
return ImportResolution(from_file=from_file, symbol=symbol)

# =====[ Check if `module/__init__.py` file exists in the graph ]=====
if from_file := self.ctx.get_file(filepath):
symbol = from_file.get_node_by_name(symbol_name)
if symbol is None:
if from_file.get_node_from_wildcard_chain(symbol_name):
Expand All @@ -163,40 +177,30 @@ def resolve_import(self, base_path: str | None = None, *, add_module_name: str |
else:
return ImportResolution(from_file=from_file, symbol=symbol)

# =====[ Check if `module/__init__.py` file exists in the graph ]=====
if from_file := self.ctx.get_file(filepath):
symbol = from_file.get_node_by_name(symbol_name)
if symbol is None:
if from_file.get_node_from_wildcard_chain(symbol_name):
return ImportResolution(from_file=from_file, symbol=None, imports_file=True)
else:
# This is most likely a broken import
return ImportResolution(from_file=from_file, symbol=None)

else:
return ImportResolution(from_file=from_file, symbol=symbol)

# =====[ Case: Can't resolve the import ]=====
if base_path == "":
# Try to resolve with "src" as the base path
return self.resolve_import(base_path="src", add_module_name=add_module_name)
if base_path == "src":
# Try "test" next
return self.resolve_import(base_path="test", add_module_name=add_module_name)
# =====[ Case: Can't resolve the import ]=====
if base_path == "":
# Try to resolve with "src" as the base path
return self.resolve_import(base_path="src", add_module_name=add_module_name)
if base_path == "src":
# Try "test" next
return self.resolve_import(base_path="test", add_module_name=add_module_name)

# if not G_override:
# for resolver in ctx.import_resolvers:
# if imp := resolver.resolve(self):
# return imp
# if not G_override:
# for resolver in ctx.import_resolvers:
# if imp := resolver.resolve(self):
# return imp

return None
# # =====[ Check if we are importing an external module in the graph ]=====
# if ext := self.ctx.get_external_module(self.source, self._unique_node.source):
# return ImportResolution(symbol=ext)
# # Implies we are not importing the symbol from the current repo.
# # In these cases, consider the import as an ExternalModule and add to graph
# ext = ExternalModule.from_import(self)
# return ImportResolution(symbol=ext)
return None
# # =====[ Check if we are importing an external module in the graph ]=====
# if ext := self.ctx.get_external_module(self.source, self._unique_node.source):
# return ImportResolution(symbol=ext)
# # Implies we are not importing the symbol from the current repo.
# # In these cases, consider the import as an ExternalModule and add to graph
# ext = ExternalModule.from_import(self)
# return ImportResolution(symbol=ext)
except AssertionError:
# Codebase is probably trying to import file from outside repo
return None

@noapidoc
@reader
Expand Down
116 changes: 60 additions & 56 deletions src/codegen/sdk/typescript/import_resolution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,63 +215,67 @@ def resolve_import(self, base_path: str | None = None, *, add_module_name: str |
- symbol: The specific symbol being imported (None for module imports)
- imports_file: True if importing the entire file/module
"""
self.file: TSFile # Type cast ts_file
base_path = base_path or self.ctx.projects[0].base_path or ""

# Get the import source path
import_source = self.module.source.strip('"').strip("'") if self.module else ""

# Try to resolve the import using the tsconfig paths
if self.file.ts_config:
import_source = self.file.ts_config.translate_import_path(import_source)

# Check if need to resolve relative import path to absolute path
relative_import = False
if import_source.startswith("."):
relative_import = True

# Insert base path
# This has the happen before the relative path resolution
if not import_source.startswith(base_path):
import_source = os.path.join(base_path, import_source)

# If the import is relative, convert it to an absolute path
if relative_import:
import_source = self._relative_to_absolute_import(import_source)
else:
import_source = os.path.normpath(import_source)

# covers the case where the import is from a directory ex: "import { postExtract } from './post'"
import_name = import_source.split("/")[-1]
if "." not in import_name:
possible_paths = ["index.ts", "index.js", "index.tsx", "index.jsx"]
for p_path in possible_paths:
if self.ctx.to_absolute(os.path.join(import_source, p_path)).exists():
import_source = os.path.join(import_source, p_path)
break

# Loop through all extensions and try to find the file
extensions = ["", ".ts", ".d.ts", ".tsx", ".d.tsx", ".js", ".jsx"]
# Try both filename with and without extension
for import_source_base in (import_source, os.path.splitext(import_source)[0]):
for extension in extensions:
import_source_ext = import_source_base + extension
if file := self.ctx.get_file(import_source_ext):
if self.is_module_import():
return ImportResolution(from_file=file, symbol=None, imports_file=True)
else:
# If the import is a named import, resolve to the named export in the file
if self.symbol_name is None:
return ImportResolution(from_file=file, symbol=None, imports_file=True)
export_symbol = file.get_export(export_name=self.symbol_name.source)
if export_symbol is None:
# If the named export is not found, it is importing a module re-export.
# In this case, resolve to the file itself and dynamically resolve the symbol later.
try:
self.file: TSFile # Type cast ts_file
base_path = base_path or self.ctx.projects[0].base_path or ""

# Get the import source path
import_source = self.module.source.strip('"').strip("'") if self.module else ""

# Try to resolve the import using the tsconfig paths
if self.file.ts_config:
import_source = self.file.ts_config.translate_import_path(import_source)

# Check if need to resolve relative import path to absolute path
relative_import = False
if import_source.startswith("."):
relative_import = True

# Insert base path
# This has the happen before the relative path resolution
if not import_source.startswith(base_path):
import_source = os.path.join(base_path, import_source)

# If the import is relative, convert it to an absolute path
if relative_import:
import_source = self._relative_to_absolute_import(import_source)
else:
import_source = os.path.normpath(import_source)

# covers the case where the import is from a directory ex: "import { postExtract } from './post'"
import_name = import_source.split("/")[-1]
if "." not in import_name:
possible_paths = ["index.ts", "index.js", "index.tsx", "index.jsx"]
for p_path in possible_paths:
if self.ctx.to_absolute(os.path.join(import_source, p_path)).exists():
import_source = os.path.join(import_source, p_path)
break

# Loop through all extensions and try to find the file
extensions = ["", ".ts", ".d.ts", ".tsx", ".d.tsx", ".js", ".jsx"]
# Try both filename with and without extension
for import_source_base in (import_source, os.path.splitext(import_source)[0]):
for extension in extensions:
import_source_ext = import_source_base + extension
if file := self.ctx.get_file(import_source_ext):
if self.is_module_import():
return ImportResolution(from_file=file, symbol=None, imports_file=True)
return ImportResolution(from_file=file, symbol=export_symbol)

# If the imported file is not found, treat it as an external module
return None
else:
# If the import is a named import, resolve to the named export in the file
if self.symbol_name is None:
return ImportResolution(from_file=file, symbol=None, imports_file=True)
export_symbol = file.get_export(export_name=self.symbol_name.source)
if export_symbol is None:
# If the named export is not found, it is importing a module re-export.
# In this case, resolve to the file itself and dynamically resolve the symbol later.
return ImportResolution(from_file=file, symbol=None, imports_file=True)
return ImportResolution(from_file=file, symbol=export_symbol)

# If the imported file is not found, treat it as an external module
return None
except AssertionError:
# Codebase is probably trying to import file from outside repo
return None

@noapidoc
@reader
Expand Down
Loading