From fb3caafd6de9fdbeb9ca2e40e8dcb6ac2e27e500 Mon Sep 17 00:00:00 2001
From: Morten Callesen <mortencallesen23@gmail.com>
Date: Tue, 15 Jun 2021 16:04:39 +0200
Subject: [PATCH 1/2] Added python_package_imports functioN

---
 d8s_python/python_data.py | 46 ++++++++++++++++++++++++++++++++++++++-
 tests/test_python_data.py | 32 +++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/d8s_python/python_data.py b/d8s_python/python_data.py
index 5ae150a..9a34a80 100644
--- a/d8s_python/python_data.py
+++ b/d8s_python/python_data.py
@@ -1,7 +1,10 @@
 import argparse
 import re
 import sys
-from typing import Any, Iterator, List, Union
+from ast import Import, ImportFrom
+from typing import Any, Dict, Iterator, List, Union
+
+from .ast_data import python_ast_objects_of_type
 
 
 # @decorators.map_firstp_arg
@@ -359,3 +362,44 @@ def python_type_name(python_type: type) -> str:
 def python_object_type_to_word(python_object: Any) -> str:
     """Convert the given python type to a string."""
     return python_type_name(type(python_object))
+
+
+def _get_importfrom_module_name(node: ImportFrom) -> str:
+    """Extract the module name from an ast.ImportFrom node.
+
+    The module name on the ast.ImportFrom node can be None for relative imports
+    In this case, this function will return the name as the dots from the import statement.
+    A few examples:
+    "from requests import get" -> "get"
+    "from . import *" -> "."
+    "from .. import *" -> ".."
+    "from .foo import bar" -> "foo"
+    """
+    if node.module is None:
+        module_name = "." * node.level
+    else:
+        module_name = node.module
+
+    return module_name
+
+
+def python_package_imports(code: str) -> Dict[str, List[str]]:
+    """Return a dictionary containing the names of all imported modules."""
+    # Start with the Import nodes.
+    # These will always have an empty list of submodules
+    # so we can just overwrite them without losing any data
+    modules = dict()
+    nodes = python_ast_objects_of_type(code, Import)
+    for node in nodes:
+        for alias in node.names:
+            modules[alias.name] = []
+
+    # Now for the ImportFrom nodes
+    importfrom_nodes = python_ast_objects_of_type(code, ImportFrom)
+    for node in importfrom_nodes:
+        module_name = _get_importfrom_module_name(node)
+
+        for alias in node.names:
+            modules.setdefault(module_name, []).append(alias.name)
+
+    return modules
diff --git a/tests/test_python_data.py b/tests/test_python_data.py
index c0e3508..f321572 100644
--- a/tests/test_python_data.py
+++ b/tests/test_python_data.py
@@ -28,6 +28,7 @@
     python_object_source_code,
     python_object_source_file,
     python_object_type_to_word,
+    python_package_imports,
     python_sort_type_list_by_name,
     python_stack_local_data,
     python_todos,
@@ -482,3 +483,34 @@ def test_python_function_blocks__async_functions():
         'def foo(n):\n    """Foo."""\n    return n',
         'async def bar(n):\n    """Some async func."""\n    return n',
     ]
+
+
+def test_python_package_imports():
+    s = ''
+    assert python_package_imports(s) == {}
+
+    s = '''import requests'''
+    assert python_package_imports(s) == {'requests': []}
+
+    s = '''from math import sqrt'''
+    assert python_package_imports(s) == {'math': ['sqrt']}
+
+    s = '''
+import os as _
+import sys,random'''
+    assert python_package_imports(s) == {'os': [], 'sys': [], 'random': []}
+
+    s = '''
+from democritus_dates import date_parse, date_now
+import requests
+
+print(date_parse('2 days from now'))
+    '''
+    assert python_package_imports(s) == {'democritus_dates': ['date_parse', 'date_now'], 'requests': []}
+
+    s = '''
+from . import everything
+from ..foo.bar import *
+    '''
+
+    assert python_package_imports(s) == {'.': ['everything'], 'foo.bar': ['*']}

From ba6ae005706a3c5be4f47b92df238614f1661abf Mon Sep 17 00:00:00 2001
From: mCallesen <1368902+mCallesen@users.noreply.github.com>
Date: Thu, 17 Jun 2021 17:08:11 +0200
Subject: [PATCH 2/2] fix documentation error in _get_importfrom_module_name

Co-authored-by: Floyd Hightower <floyd.hightower27@gmail.com>
---
 d8s_python/python_data.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/d8s_python/python_data.py b/d8s_python/python_data.py
index 9a34a80..cab32ee 100644
--- a/d8s_python/python_data.py
+++ b/d8s_python/python_data.py
@@ -370,7 +370,7 @@ def _get_importfrom_module_name(node: ImportFrom) -> str:
     The module name on the ast.ImportFrom node can be None for relative imports
     In this case, this function will return the name as the dots from the import statement.
     A few examples:
-    "from requests import get" -> "get"
+    "from requests import get" -> "requests"
     "from . import *" -> "."
     "from .. import *" -> ".."
     "from .foo import bar" -> "foo"