Improve module loading

di · di · commit c96b89fdee09 · 2020-05-19T19:08:56.000-05:00
diff --git a/src/functions_framework/__init__.py b/src/functions_framework/__init__.py
@@ -140,10 +140,24 @@ def create_app(target=None, source=None, signature_type=None):
     # Set the environment variable if it wasn't already
     os.environ["FUNCTION_SIGNATURE_TYPE"] = signature_type
 
-    # Load the source file
-    spec = importlib.util.spec_from_file_location("__main__", source)
+    # Load the source file:
+    # 1. Extract the module name from the source path
+    realpath = os.path.realpath(source)
+    directory, filename = os.path.split(realpath)
+    name, extension = os.path.splitext(filename)
+
+    # 2. Create a new module
+    spec = importlib.util.spec_from_file_location(name, realpath)
     source_module = importlib.util.module_from_spec(spec)
-    sys.path.append(os.path.dirname(os.path.realpath(source)))
+
+    # 3. Add the directory of the source to sys.path to allow the function to
+    # load modules relative to its location
+    sys.path.append(directory)
+
+    # 4. Add the module to sys.modules
+    sys.modules[name] = source_module
+
+    # 5. Execute the module
     spec.loader.exec_module(source_module)
 
     app = flask.Flask(target, template_folder=template_folder)
diff --git a/tests/test_functions/module_is_correct/main.py b/tests/test_functions/module_is_correct/main.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os.path
 import typing
 
 
@@ -21,7 +22,9 @@ class TestClass:
 
 def function(request):
     # Ensure that the module for any object in this file is set correctly
-    assert TestClass.__mro__[0].__module__ == "__main__"
+    _, filename = os.path.split(__file__)
+    name, _ = os.path.splitext(filename)
+    assert TestClass.__mro__[0].__module__ == name
 
     # Ensure that calling `get_type_hints` on an object in this file succeeds
     assert typing.get_type_hints(TestClass) == {}
