1"""Handling of module and package related details."""
2
3import dataclasses
4from typing import Any
5
6from pytype import file_utils
7from pytype import module_utils
8from pytype.pyi.types import ParseError  # pylint: disable=g-importing-member
9from pytype.pytd import pytd
10from pytype.pytd import visitors
11from pytype.pytd.parse import parser_constants  # pylint: disable=g-importing-member
12
13
14@dataclasses.dataclass
15class Import:
16  """Result of processing an import statement."""
17
18  pytd_node: Any
19  name: str
20  new_name: str
21  qualified_name: str = ""
22
23  def pytd_alias(self):
24    return pytd.Alias(self.new_name, self.pytd_node)
25
26
27class Module:
28  """Module and package details."""
29
30  def __init__(self, filename, module_name):
31    self.filename = filename
32    self.module_name = module_name
33    is_package = file_utils.is_pyi_directory_init(filename)
34    self.package_name = module_utils.get_package_name(module_name, is_package)
35    self.parent_name = module_utils.get_package_name(self.package_name, False)
36
37  def _qualify_name_with_special_dir(self, orig_name):
38    """Handle the case of '.' and '..' as package names."""
39    if "__PACKAGE__." in orig_name:
40      # Generated from "from . import foo" - see parser.yy
41      prefix, _, name = orig_name.partition("__PACKAGE__.")
42      if prefix:
43        raise ParseError(f"Cannot resolve import: {orig_name}")
44      return f"{self.package_name}.{name}"
45    elif "__PARENT__." in orig_name:
46      # Generated from "from .. import foo" - see parser.yy
47      prefix, _, name = orig_name.partition("__PARENT__.")
48      if prefix:
49        raise ParseError(f"Cannot resolve import: {orig_name}")
50      if not self.parent_name:
51        raise ParseError(
52            f"Cannot resolve relative import ..: Package {self.package_name} "
53            "has no parent"
54        )
55      return f"{self.parent_name}.{name}"
56    else:
57      return None
58
59  def qualify_name(self, orig_name):
60    """Qualify an import name."""
61    # Doing the "builtins" rename here ensures that we catch alias names.
62    orig_name = visitors.RenameBuiltinsPrefixInName(orig_name)
63    if not self.package_name:
64      return orig_name
65    rel_name = self._qualify_name_with_special_dir(orig_name)
66    if rel_name:
67      return rel_name
68    if orig_name.startswith("."):
69      name = module_utils.get_absolute_name(self.package_name, orig_name)
70      if name is None:
71        raise ParseError(
72            f"Cannot resolve relative import {orig_name.rsplit('.', 1)[0]}")
73      return name
74    return orig_name
75
76  def process_import(self, item):
77    """Process 'import a, b as c, ...'."""
78    if not isinstance(item, tuple):
79      # We don't care about imports that are not aliased.
80      return None
81    name, new_name = item
82    module_name = self.qualify_name(name)
83    as_name = self.qualify_name(new_name)
84    t = pytd.Module(name=as_name, module_name=module_name)
85    return Import(pytd_node=t, name=name, new_name=new_name)
86
87  def process_from_import(self, from_package, item):
88    """Process 'from a.b.c import d, ...'."""
89    if isinstance(item, tuple):
90      name, new_name = item
91    else:
92      name = new_name = item
93    qualified_name = self.qualify_name(f"{from_package}.{name}")
94    # We should ideally not need this check, but we have typing
95    # special-cased in some places.
96    if not qualified_name.startswith("typing.") and name != "*":
97      # Mark this as an externally imported type, so that AddNamePrefix
98      # does not prefix it with the current package name.
99      qualified_name = (parser_constants.EXTERNAL_NAME_PREFIX +
100                        qualified_name)
101    t = pytd.NamedType(qualified_name)
102    if name == "*":
103      # A star import is stored as
104      # 'imported_mod.* = imported_mod.*'. The imported module needs to be
105      # in the alias name so that multiple star imports are handled
106      # properly. LookupExternalTypes() replaces the alias with the
107      # contents of the imported module.
108      assert new_name == name
109      new_name = t.name
110    return Import(pytd_node=t, name=name, new_name=new_name,
111                  qualified_name=qualified_name)
112