diff --git a/armi/utils/codeTiming.py b/armi/utils/codeTiming.py index 543b00132..b205bbe11 100644 --- a/armi/utils/codeTiming.py +++ b/armi/utils/codeTiming.py @@ -16,6 +16,7 @@ import copy import os import time +import functools def timed(*args): @@ -36,9 +37,7 @@ def mymethod2(stuff) """ def time_decorator(func): - time_decorator.__doc__ = func.__doc__ - time_decorator.__name__ = func.__name__ - + @functools.wraps(func) def time_wrapper(*args, **kwargs): generated_name = "::".join( [ diff --git a/armi/utils/tests/test_utils.py b/armi/utils/tests/test_utils.py index d5aac0314..19a221ab5 100644 --- a/armi/utils/tests/test_utils.py +++ b/armi/utils/tests/test_utils.py @@ -36,6 +36,7 @@ getPreviousTimeNode, getCumulativeNodeNum, hasBurnup, + codeTiming, ) @@ -154,6 +155,19 @@ def test_classesInHierarchy(self): self.assertGreater(len(r.core.getAssemblies()), 50) self.assertGreater(len(r.core.getBlocks()), 200) + def test_codeTiming(self): + """ + Test that codeTiming preserves function attributes when it wraps a function + """ + + @codeTiming.timed + def testFunc(): + """Test function docstring""" + pass + + self.assertEqual(getattr(testFunc, "__doc__"), "Test function docstring") + self.assertEqual(getattr(testFunc, "__name__"), "testFunc") + class CyclesSettingsTests(unittest.TestCase): """ diff --git a/doc/release/0.2.rst b/doc/release/0.2.rst index 883418687..631fd05a8 100644 --- a/doc/release/0.2.rst +++ b/doc/release/0.2.rst @@ -13,6 +13,7 @@ What's new in ARMI #. Removed all old ARMI requirements, to start the work fresh. (`PR#1438 `_) #. Downgrading Draft PRs as policy. (`PR#1444 `_) #. Attempt to set representative block number densities by component if possible. (`PR#1412 `_) +#. Use functools to preserve function attributes when wrapping with codeTiming.timed (`PR#1466 `_) #. TBD Bug fixes