From 5e671b39db30a4cbf63a218aed4abc02acce288f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= <pawamoy@pm.me>
Date: Fri, 20 Aug 2021 17:38:38 +0200
Subject: [PATCH] fix: Fix getting parent module of decorated functions

Issue mkdocstrings#162: https://github.com/mkdocstrings/mkdocstrings/issues/162
---
 src/pytkdocs/loader.py               |  4 ++--
 tests/fixtures/decorated_function.py | 16 ++++++++++++++++
 tests/test_loader.py                 | 11 +++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 tests/fixtures/decorated_function.py

diff --git a/src/pytkdocs/loader.py b/src/pytkdocs/loader.py
index fe60edf..2f29467 100644
--- a/src/pytkdocs/loader.py
+++ b/src/pytkdocs/loader.py
@@ -434,9 +434,9 @@ def get_module_documentation(self, node: ObjectNode, select_members=None) -> Mod
         for member_name, member in inspect.getmembers(module):
             if self.select(member_name, select_members):  # type: ignore
                 child_node = ObjectNode(member, member_name, parent=node)
-                if child_node.is_class() and node.root.obj is inspect.getmodule(member):
+                if child_node.is_class() and node.root.obj is inspect.getmodule(child_node.obj):
                     root_object.add_child(self.get_class_documentation(child_node))
-                elif child_node.is_function() and node.root.obj is inspect.getmodule(member):
+                elif child_node.is_function() and node.root.obj is inspect.getmodule(child_node.obj):
                     root_object.add_child(self.get_function_documentation(child_node))
                 elif member_name in attributes_data:
                     root_object.add_child(self.get_attribute_documentation(child_node, attributes_data[member_name]))
diff --git a/tests/fixtures/decorated_function.py b/tests/fixtures/decorated_function.py
new file mode 100644
index 0000000..2e72dff
--- /dev/null
+++ b/tests/fixtures/decorated_function.py
@@ -0,0 +1,16 @@
+from functools import lru_cache
+
+
+@lru_cache()
+def add(a, b):
+    return a + b
+
+
+# control
+def sub(a, b):
+    return a - b
+
+
+# simulating a decorator that does not set __module__ properly
+# on the wrapper object
+del add.__module__
diff --git a/tests/test_loader.py b/tests/test_loader.py
index 845cce4..4af308d 100644
--- a/tests/test_loader.py
+++ b/tests/test_loader.py
@@ -534,3 +534,14 @@ def test_method_descriptor():
     assert len(obj.signature.parameters) == 2
     assert obj.docstring
     assert obj.category == "method"
+
+
+def test_load_decorated_function():
+    """Load a decorated function."""
+    loader = Loader(new_path_syntax=True)
+    obj = loader.get_object_documentation("tests.fixtures.decorated_function")
+    assert [child.name for child in obj.children] == ["add", "sub"]
+    for child in obj.children:
+        assert child.category == "function"
+        assert child.parent is child.root
+        assert child.parent.name == "decorated_function"