diff --git a/benchmarks/benchmarks.py b/benchmarks/benchmarks.py index 9e52c78..49c27cb 100644 --- a/benchmarks/benchmarks.py +++ b/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 diff --git a/docs/notebooks/timing.ipynb b/docs/notebooks/timing.ipynb index 78a8a83..5988b41 100644 --- a/docs/notebooks/timing.ipynb +++ b/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, diff --git a/src/citation_compass/citation.py b/src/citation_compass/citation.py index dde46d3..6fc5813 100644 --- a/src/citation_compass/citation.py +++ b/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): diff --git a/tests/citation_compass/test_citation.py b/tests/citation_compass/test_citation.py index a23cb59..a8f8cfe 100644 --- a/tests/citation_compass/test_citation.py +++ b/tests/citation_compass/test_citation.py @@ -1,5 +1,6 @@ import fake_module import pytest +import types from citation_compass import ( cite_function, @@ -31,6 +32,34 @@ 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 = [ @@ -38,6 +67,9 @@ def test_citations_all(): "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()