Fix #5 and fit nameof in more cases

Author: pwwang

README.md

@@ -210,16 +210,16 @@ a == b
 
 ## Limitations
 - Working in `ipython REPL` but not in standard `python console`
-- You have to know at which stack the function/class will be called
+- You have to know at which stack the function/class will be called (caller's depth)
 - Not working with `reticulate` from `R` since it cuts stacks to the most recent one.
-- `nameof` cannot be used in statements in `pytest`
-  ```
-  a = 1
-  assert nameof(a) == 'a'
-  # Retrieving failure.
-  # The right way:
-  aname = nameof(a)
-  assert aname == 'a'
+- ~~`nameof` cannot be used in statements in `pytest`~~ (supported in `v0.2.0`)
+  ```diff
+  -a = 1
+  +assert nameof(a) == 'a'
+  -# Retrieving failure.
+  -# The right way:
+  -aname = nameof(a)
+  -assert aname == 'a'
   ```
 
 [1]: https://github.com/pwwang/python-varname

tests/test_varname.py

@@ -7,7 +7,7 @@
                      will,
                      inject,
                      namedtuple,
-                     _get_executing,
+                     _get_node,
                      nameof)
 
 @pytest.fixture
@@ -246,8 +246,9 @@ def func():
     f = func()
     assert f == 'fabc'
 
-    with pytest.raises(VarnameRetrievingError):
-        assert nameof(f)
+    assert nameof(f) == 'f'
+    assert 'f' == nameof(f)
+    assert len(nameof(f)) == 1
 
     fname1 = fname = nameof(f)
     assert fname1 == fname == 'f'
@@ -258,6 +259,44 @@ def func():
     with pytest.raises(VarnameRetrievingError):
         nameof()
 
+def test_nameof_statements():
+    a = {'test': 1}
+    test = {}
+    del a[nameof(test)]
+    assert a == {}
+
+    def func():
+        return nameof(test)
+
+    assert func() == 'test'
+
+    def func2():
+        yield nameof(test)
+
+    assert list(func2()) == ['test']
+
+    def func3():
+        raise ValueError(nameof(test))
+
+    with pytest.raises(ValueError) as verr:
+        func3()
+    assert str(verr.value) == 'test'
+
+    for i in [0]:
+        assert nameof(test) == 'test'
+        assert len(nameof(test)) == 4
+
+def test_nameof_expr():
+    import varname
+    test = {}
+    assert len(varname.nameof(test)) == 4
+
+    lam = lambda: 0
+    lam.a = 1
+    with pytest.raises(VarnameRetrievingError) as vrerr:
+        varname.nameof(test, lam.a)
+    assert str(vrerr.value) == ("Only variables should "
+                                "be passed to nameof.")
 
 def test_class_property():
     class C:
@@ -327,7 +366,7 @@ def get_will():
 def test_frame_fail(monkey_patch):
     """Test when failed to retrieve the frame"""
     # Let's monkey-patch inspect.stack to do this
-    assert _get_executing(1) is None
+    assert _get_node(1) is None
 
 def test_frame_fail_varname(monkey_patch):
 

varname.py

@@ -18,18 +18,83 @@ class VarnameRetrievingWarning(Warning):
 class VarnameRetrievingError(Exception):
     """When failed to retrieve the varname"""
 
-def _get_executing(caller):
-    """Try to get the executing object
+def _get_node(caller):
+    """Try to get node from the executing object.
+
     This can fail when a frame is failed to retrieve.
     One case should be when python code is executed in
     R pacakge `reticulate`, where only first frame is kept.
+
+    When the node can not be retrieved, try to return the first statement.
     """
     try:
         frame = inspect.stack()[caller+2].frame
     except IndexError:
         return None
     else:
-        return executing.Source.executing(frame)
+        exet = executing.Source.executing(frame)
+
+    if exet.node:
+        return exet.node
+
+    return list(exet.statements)[0]
+
+def _lookfor_parent_assign(node):
+    """Look for an ast.Assign node in the parents"""
+    while hasattr(node, 'parent'):
+        node = node.parent
+
+        if isinstance(node, ast.Assign):
+            return node
+    return None
+
+def _lookfor_child_nameof(node):
+    """Look for ast.Call with func=Name(id='nameof',...)"""
+    # pylint: disable=too-many-return-statements
+    if isinstance(node, ast.Call):
+        # if node.func.id == 'nameof':
+        #     return node
+
+        # We want to support alias for nameof, i.e. nameof2
+        # Or called like: varname.nameof(test)
+        # If all args are ast.Name, then if must be alias of nameof
+        # Since this is originated from it, and there is no other
+        # ast.Call node in args
+        if not any(isinstance(arg, ast.Call) for arg in node.args):
+            return node
+
+        # print(nameof(test))
+        for arg in node.args:
+            found = _lookfor_child_nameof(arg)
+            if found:
+                return found
+
+    elif isinstance(node, ast.Compare):
+        # nameof(test) == 'test'
+        found = _lookfor_child_nameof(node.left)
+        if found:
+            return found
+        # 'test' == nameof(test)
+        for comp in node.comparators:
+            found = _lookfor_child_nameof(comp)
+            if found:
+                return found
+
+    elif isinstance(node, ast.Assert):
+        # assert nameof(test) == 'test'
+        found = _lookfor_child_nameof(node.test)
+        if found:
+            return found
+
+    elif isinstance(node, ast.Expr): # pragma: no cover
+        # print(nameof(test)) in ipython's forloop
+        # issue #5
+        found = _lookfor_child_nameof(node.value)
+        if found:
+            return found
+
+    return None
+
 
 def varname(caller=1, raise_exc=False):
     """Get the variable name that assigned by function/class calls
@@ -53,42 +118,40 @@ def varname(caller=1, raise_exc=False):
             in the assign node. (e.g: `a = b = func()`, in such a case,
             `b == 'a'`, may not be the case you want)
     """
-    exec_obj = _get_executing(caller)
-    if not exec_obj:
+    node = _get_node(caller)
+    if not node:
         if raise_exc:
-            raise VarnameRetrievingError("Unable to retrieve the frame.")
+            raise VarnameRetrievingError("Unable to retrieve the ast node.")
         VARNAME_INDEX[0] += 1
         warnings.warn(f"var_{VARNAME_INDEX[0]} used.",
                       VarnameRetrievingWarning)
         return f"var_{VARNAME_INDEX[0]}"
 
-    node = exec_obj.node
-    while hasattr(node, 'parent'):
-        node = node.parent
-
-        if isinstance(node, ast.Assign):
-            # Need to actually check that there's just one
-            if len(node.targets) > 1:
-                warnings.warn("Multiple targets in assignment, variable name "
-                              "on the very left will be used.",
-                              MultipleTargetAssignmentWarning)
-            target = node.targets[0]
-
-            # Need to check that it's a variable
-            if isinstance(target, ast.Name):
-                return target.id
-
+    node = _lookfor_parent_assign(node)
+    if not node:
+        if raise_exc:
             raise VarnameRetrievingError(
-                f"Invaid variable assigned: {ast.dump(target)!r}"
+                'Failed to retrieve the variable name.'
             )
+        VARNAME_INDEX[0] += 1
+        warnings.warn(f"var_{VARNAME_INDEX[0]} used.", VarnameRetrievingWarning)
+        return f"var_{VARNAME_INDEX[0]}"
+
+    # Need to actually check that there's just one
+    # give warnings if: a = b = func()
+    if len(node.targets) > 1:
+        warnings.warn("Multiple targets in assignment, variable name "
+                      "on the very left will be used.",
+                      MultipleTargetAssignmentWarning)
+    target = node.targets[0]
 
-    if raise_exc:
-        raise VarnameRetrievingError('Failed to retrieve the variable name.')
+    # must be a variable
+    if isinstance(target, ast.Name):
+        return target.id
 
-    VARNAME_INDEX[0] += 1
-    warnings.warn(f"var_{VARNAME_INDEX[0]} used.",
-                  VarnameRetrievingWarning)
-    return f"var_{VARNAME_INDEX[0]}"
+    raise VarnameRetrievingError(
+        f"Invaid variable assigned: {ast.dump(target)!r}"
+    )
 
 def will(caller=1, raise_exc=False):
     """Detect the attribute name right immediately after a function call.
@@ -118,15 +181,14 @@ def i_will():
         VarnameRetrievingError: When `raise_exc` is `True` and we failed to
             detect the attribute name (including not having one)
     """
-    exec_obj = _get_executing(caller)
-    if not exec_obj:
+    node = _get_node(caller)
+    if not node:
         if raise_exc:
             raise VarnameRetrievingError("Unable to retrieve the frame.")
         return None
 
     ret = None
     try:
-        node = exec_obj.node
         # have to be used in a call
         assert isinstance(node, (ast.Attribute, ast.Call)), (
             "Invalid use of function `will`"
@@ -174,24 +236,19 @@ def nameof(*args):
         *args: A couple of variables passed in
 
     Returns:
-        tuple|str:
+        tuple|str: The names of variables passed in
     """
-    exec_obj = _get_executing(0)
-    if not exec_obj:
-        raise VarnameRetrievingError("Unable to retrieve the frame.")
-
-    if not exec_obj.node:
-        # we cannot do: assert nameof(a) == 'a' in pytest
-        raise VarnameRetrievingError("Callee's node cannot be detected.")
-
-    assert isinstance(exec_obj.node, ast.Call)
+    node = _get_node(0)
+    node = _lookfor_child_nameof(node)
+    if not node:
+        raise VarnameRetrievingError("Unable to retrieve callee's node.")
 
     ret = []
-    for node in exec_obj.node.args:
-        if not isinstance(node, ast.Name):
+    for arg in node.args:
+        if not isinstance(arg, ast.Name):
             raise VarnameRetrievingError("Only variables should "
                                          "be passed to nameof.")
-        ret.append(node.id)
+        ret.append(arg.id)
 
     if not ret:
         raise VarnameRetrievingError("At least one variable should be "