Repository: fniessink/access-modifiers

Issue:
Make access modifiers work when called from a list comprehension

Comments:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a780331..e14d12f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,12 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 <!-- The line "## <square-bracket>Unreleased</square-bracket>" is replaced by the ci/release.py script with the new release version and release date. -->
 
+## [Unreleased]
+
+### Fixed
+
+- Make the access modifiers work when called from a dict comprehension, list comprehension, lambda, or other code block that gets its own frame. Fixes [#5](https://github.com/fniessink/access-modifiers/issues/5).
+
 ## [0.1.3] - [2019-08-24]
 
 ### Fixed
diff --git a/access_modifiers/access_modifiers.py b/access_modifiers/access_modifiers.py
index 948c72c..b8a56d9 100644
--- a/access_modifiers/access_modifiers.py
+++ b/access_modifiers/access_modifiers.py
@@ -19,11 +19,15 @@ def privatemethod(method: Callable[..., ReturnType]) -> Callable[..., ReturnType
     def private_method_wrapper(*args, **kwargs) -> ReturnType:
         """Wrap the original method to make it private."""
         caller_frame = getframe(1)
+        caller_code = caller_frame.f_code
+        caller_name = caller_code.co_name
+        while caller_name.startswith("<"):  # Code is a <lambda>, <dictcomp>, <listcomp>, or other non-method code block
+            caller_frame = caller_frame.f_back
+            caller_code = caller_frame.f_code
+            caller_name = caller_code.co_name
         caller_instance = caller_frame.f_locals.get("self")
         if caller_instance is not args[0]:
             raise AccessException(f"Attempted call to private method {method} from another object")
-        caller_code = caller_frame.f_code
-        caller_name = caller_code.co_name
         # Look up the calling method to see if it's defined in the same class as the private method
         for caller_class in caller_instance.__class__.mro():
             caller = caller_class.__dict__.get(caller_name)


Test patch:
diff --git a/access_modifiers/tests/test_private.py b/access_modifiers/tests/test_private.py
index 1d113e6..bf2aff1 100644
--- a/access_modifiers/tests/test_private.py
+++ b/access_modifiers/tests/test_private.py
@@ -56,6 +56,51 @@ class PrivateMethodTest(unittest.TestCase):
 
         self.assertRaises(AccessException, Subclass().public_method)
 
+    def test_call_private_method_via_list_comprehension(self):
+        """Test that accessing a private method via a list comprehension in a public method works."""
+
+        class Class:
+            @privatemethod
+            def private_method(self):  # pylint: disable=no-self-use
+                return "Class.private_method"
+
+            def public_method(self):
+                return ["Class.public_method -> " + self.private_method() for _ in range(1)][0]
+
+        self.assertEqual("Class.public_method -> Class.private_method", Class().public_method())
+
+    def test_call_private_method_via_nested_lambda(self):
+        """Test that accessing a private method via a nested lambda in a public method works."""
+
+        class Class:
+            @privatemethod
+            def private_method(self):  # pylint: disable=no-self-use
+                return "Class.private_method"
+
+            def public_method(self):
+                # pylint: disable=unnecessary-lambda
+                inner_lambda_function = lambda: self.private_method()
+                outer_lambda_function = lambda: "Class.public_method -> " + inner_lambda_function()
+                return outer_lambda_function()
+
+        self.assertEqual("Class.public_method -> Class.private_method", Class().public_method())
+
+    def test_call_private_method_from_try_except_block(self):
+        """Test that accessing a private method from a try/except in a public method works."""
+
+        class Class:
+            @privatemethod
+            def private_method(self):  # pylint: disable=no-self-use
+                return "Class.private_method"
+
+            def public_method(self):
+                try:
+                    return "Class.public_method -> " + self.private_method()
+                except AttributeError:  # pragma: nocover
+                    pass
+
+        self.assertEqual("Class.public_method -> Class.private_method", Class().public_method())
+
     def test_override_private_method(self):
         """Test that an overridden private method can't call its super."""