Fix the types for cite_function

benchmarks/benchmarks.py

@@ -6,13 +6,13 @@
 import citation_compass as cc
 
 
-@cc.cite_function("fake")
+@cc.cite_function(label="fake")
 def fake_function():
     """A fake function to demonstrate the use of the citation_compass package."""
     return 1
 
 
-@cc.cite_function("fake", track_used=False)
+@cc.cite_function(label="fake", track_used=False)
 def fake_function2():
     """A fake function to demonstrate the use of the citation_compass package."""
     return 1
@@ -28,7 +28,7 @@ def __init__(self):
 def time_create_function():
     """Time the use of a wrapper with a label."""
 
-    @cc.cite_function("example")
+    @cc.cite_function(label="example")
     def test_function():
         return 1
 

docs/notebooks/timing.ipynb

@@ -25,13 +25,13 @@
     "    return 1\n",
     "\n",
     "\n",
-    "@cc.cite_function(\"test_func2\")\n",
+    "@cc.cite_function(label=\"test_func2\")\n",
     "def test_func2():\n",
     "    \"\"\"A test function with a citation\"\"\"\n",
     "    return 1\n",
     "\n",
     "\n",
-    "@cc.cite_function(\"test_func3\", track_used=False)\n",
+    "@cc.cite_function(label=\"test_func3\", track_used=False)\n",
     "def test_func3():\n",
     "    \"\"\"A test function with a citation that will not be tracked\"\"\"\n",
     "    return 1"
@@ -94,7 +94,7 @@
     "        return (-b + math.sqrt(inner)) / (2 * a), (-b - math.sqrt(inner)) / (2 * a)\n",
     "\n",
     "\n",
-    "@cc.cite_function(\"test_func4\")\n",
+    "@cc.cite_function\n",
     "def test_func4(a, b, c):\n",
     "    \"\"\"A test function with a citation\"\"\"\n",
     "    inner = b**2 - 4 * a * c\n",
@@ -127,7 +127,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "citation",
+   "display_name": "citation (3.13.8)",
    "language": "python",
    "name": "python3"
   },
@@ -141,7 +141,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.13.8"
   }
  },
  "nbformat": 4,

src/citation_compass/citation.py

@@ -3,6 +3,7 @@
 from functools import wraps
 from os import urandom
 import sys
+import types
 
 from citation_compass.citation_registry import (
     CitationEntry,
@@ -74,14 +75,18 @@ def init_wrapper(*args, **kwargs):
         cls.__init__ = init_wrapper
 
 
-def cite_function(label=None, track_used=True):
+def cite_function(callable=None, *, label=None, track_used=True):
     """A function wrapper for adding a citation to a function or
     class method.
 
     Parameters
     ----------
+    callable : function or method, optional
+        The function or method to add a citation to. This is automatically passed as
+        the first argument when using the decorator without parentheses.
     label : str, optional
-        The (optional) user-defined label for the citation.
+        The (optional) user-defined label for the citation. If not provided,
+        the label will be auto-extracted from the function's docstring.
     track_used : bool
         If True, the function will be marked as used when it is called.
         This adds a small amount of overhead to each function call.
@@ -92,35 +97,47 @@ class method.
     function
         The wrapped function or method.
     """
-    # If the label is callable, there were no parentheses on the
-    # dectorator and it passed in the function instead. So use None
-    # as the label.
-    use_label = label if not callable(label) else None
-
-    def decorator(func):
-        entry = CitationEntry.from_object(func, label=use_label)
-        CITATION_COMPASS_REGISTRY.add(entry)
+    # This decorator is designed as a two-layer decorator. The first layer handles the (optional)
+    # arguments. The second handles the actual function wrapping.
 
-        # Wrap the function so it is marked as USED when it is called.
-        if track_used:
+    def _inner_decorator(callable):
+        # The inner decorator is used to set the "all" citations entry and handle
+        # the correct return types when wrapping the function.
 
-            @wraps(func)
-            def fun_wrapper(*args, **kwargs):
-                # Save the citation as USED when it is first called.
-                CITATION_COMPASS_REGISTRY.mark_used(entry.key)
-                return func(*args, **kwargs)
-        else:
-            # We do not wrap the function, but just return the original function.
-            fun_wrapper = func
+        # Add the function to the registry (for the "all" citations).
+        entry = CitationEntry.from_object(callable, label=label)
+        CITATION_COMPASS_REGISTRY.add(entry)
 
-            # We mark as used be default so the citation does not get dropped.
+        # If we are not tracking when the function is used, we don't need to wrap it.
+        # We can just return the original callable.
+        if not track_used:
+            # We mark as used by default so the citation does not get dropped.
             CITATION_COMPASS_REGISTRY.mark_used(entry.key)
+            return callable
 
-        return fun_wrapper
+        # If the callable is a classmethod or method, we need to get the function that has
+        # the self or cls argument.
+        func = callable.__func__ if isinstance(callable, (classmethod, types.MethodType)) else callable
 
-    if callable(label):
-        return decorator(label)
-    return decorator
+        # Define the actual wrapper for the callable we passed in. This wrapper function will
+        # be called each time the internal function is called.
+        @wraps(func)
+        def citation_wrapper(*args, **kwargs):
+            # Save the citation as USED when it is first called.
+            CITATION_COMPASS_REGISTRY.mark_used(entry.key)
+            return func(*args, **kwargs)
+
+        # We cast the wrapped function as the correct type.
+        if isinstance(callable, classmethod):
+            return classmethod(citation_wrapper)
+        elif isinstance(callable, staticmethod):
+            return staticmethod(citation_wrapper)
+        elif isinstance(callable, types.MethodType):
+            return types.MethodType(citation_wrapper, callable.__self__)
+        return citation_wrapper
+
+    # Handle the optional parentheses in the decorator.
+    return _inner_decorator if callable is None else _inner_decorator(callable)
 
 
 def cite_object(obj, label=None):

tests/citation_compass/test_citation.py

@@ -1,5 +1,6 @@
 import fake_module
 import pytest
+import types
 
 from citation_compass import (
     cite_function,
@@ -31,13 +32,44 @@ def example_function_x(x):
     return x
 
 
+class _FakeTestingClass:
+    """A fake class for testing."""
+
+    def __init__(self, data=1):
+        self.data = data
+
+    @classmethod
+    @cite_function
+    def fake_class_classmethod(cls):
+        """A fake classmethod for testing."""
+        return cls(data=1)
+
+    @staticmethod
+    @cite_function
+    def fake_class_staticmethod():
+        """A fake staticmethod for testing."""
+        return 0
+
+    @cite_function
+    def fake_class_normal_method(self):
+        """A fake normal class method for testing."""
+        return self.data
+
+    def uncited_method(self):
+        """A method that is not cited."""
+        return self.data
+
+
 def test_citations_all():
     """Check that all the citations are registered."""
     known_citations = [
         # The functions defined in this file.
         "test_citation.example_function_1: function_citation_1",
         "test_citation.example_function_2: function_citation_2",
         "test_citation.example_function_x: function_citation_x",
+        "test_citation._FakeTestingClass.fake_class_classmethod: A fake classmethod for testing.",
+        "test_citation._FakeTestingClass.fake_class_staticmethod: A fake staticmethod for testing.",
+        "test_citation._FakeTestingClass.fake_class_normal_method: A fake normal class method for testing.",
         # The items defined in fake_module.
         "fake_module: CitationCompass, 2025.",
         "fake_module.FakeClass.fake_method: A fake class method for testing.",
@@ -60,8 +92,8 @@ def test_citations_all():
     obj = fake_module.FakeCitedClass()
     assert isinstance(obj, fake_module.FakeCitedClass)
 
-    # A citation with no docstring, but a label.
-    @cite_function("function_citation_3")
+    # A citation with no docstring, but a manual label.
+    @cite_function(label="function_citation_3")
     def example_function_3():
         return 3
 
@@ -238,6 +270,50 @@ def test_find_in_citations():
     assert len(find_in_citations("FakeCitedClass", True)) == 1
 
 
+def test_functions_in_class():
+    """Test that we correctly handle methods in a class including static and class methods."""
+    obj = _FakeTestingClass(data=5)
+
+    # Nothing is used.
+    assert len(find_in_citations("fake_class_normal_method", True)) == 0
+    assert len(find_in_citations("fake_class_staticmethod", True)) == 0
+    assert len(find_in_citations("fake_class_classmethod", True)) == 0
+
+    # All the functions are usable.
+    assert obj.fake_class_normal_method() == 5
+    assert obj.fake_class_staticmethod() == 0
+
+    obj2 = obj.fake_class_classmethod()
+    assert isinstance(obj2, _FakeTestingClass)
+    assert obj2.data == 1
+
+    # Everything is now cited.
+    assert len(find_in_citations("fake_class_normal_method", True)) == 1
+    assert len(find_in_citations("fake_class_staticmethod", True)) == 1
+    assert len(find_in_citations("fake_class_classmethod", True)) == 1
+
+    # We preserve the types of each method when called from the class. The class method
+    # static method should be those types.
+    assert isinstance(_FakeTestingClass.__dict__["fake_class_normal_method"], types.FunctionType)
+    assert isinstance(_FakeTestingClass.__dict__["fake_class_staticmethod"], staticmethod)
+    assert isinstance(_FakeTestingClass.__dict__["fake_class_classmethod"], classmethod)
+
+    # Check the types when accessing an instance.
+    assert isinstance(obj.fake_class_normal_method, types.MethodType)
+    assert isinstance(obj.fake_class_staticmethod, types.FunctionType)
+    assert isinstance(obj.fake_class_classmethod, types.MethodType)
+
+    # We preserve the names of the methods.
+    assert obj.fake_class_classmethod.__name__ == "fake_class_classmethod"
+    assert obj.fake_class_staticmethod.__name__ == "fake_class_staticmethod"
+    assert obj.fake_class_normal_method.__name__ == "fake_class_normal_method"
+
+    # We preserve the docstring of the methods.
+    assert obj.fake_class_classmethod.__doc__ == "A fake classmethod for testing."
+    assert obj.fake_class_staticmethod.__doc__ == "A fake staticmethod for testing."
+    assert obj.fake_class_normal_method.__doc__ == "A fake normal class method for testing."
+
+
 def test_citation_context():
     """Test the CitationContext class."""
     reset_used_citations()