Merge pull request #19 from lukeyeager/redecorate

Allow a function to be decorated twice

Author: lukeyeager

flask_autodoc/autodoc.py

@@ -61,7 +61,7 @@ def nl2br(eval_ctx, value):
                                  for p in _paragraph_re.split(value))
             return result
 
-    def doc(self, groups=None, **properties):
+    def doc(self, groups=None, set_location=True, **properties):
         """Add flask route to autodoc for automatic documentation
 
         Any route decorated with this method will be added to the list of
@@ -71,32 +71,41 @@ def doc(self, groups=None, **properties):
         By specifying group or groups argument, the route can be added to one
         or multiple other groups as well, besides the 'all' group.
 
+        If set_location is True, the location of the function will be stored.
+        NOTE: this assumes that the decorator is placed just before the
+        function (in the normal way).
+
         Custom parameters may also be passed in beyond groups, if they are
-        named something not already in the dict descibed in the docstring for 
+        named something not already in the dict descibed in the docstring for
         the generare() function, they will be added to the route's properties,
         which can be accessed from the template.
 
         If a parameter is passed in with a name that is already in the dict, but
         not of a reserved name, the passed parameter overrides that dict value.
         """
         def decorator(f):
-            # Set group[s]
-            if type(groups) is list:
-                groupset = set(groups)
+            # Get previous group list (if any)
+            if f in self.func_groups:
+                groupset = self.func_groups[f]
             else:
                 groupset = set()
-                if type(groups) is str:
-                    groupset.add(groups)
+
+            # Set group[s]
+            if type(groups) is list:
+                groupset.update(groups)
+            elif type(groups) is str:
+                groupset.add(groups)
             groupset.add('all')
             self.func_groups[f] = groupset
             self.func_props[f] = properties
 
             # Set location
-            caller_frame = inspect.stack()[1]
-            self.func_locations[f] = {
-                    'filename': caller_frame[1],
-                    'line':     caller_frame[2],
-                    }
+            if set_location:
+                caller_frame = inspect.stack()[1]
+                self.func_locations[f] = {
+                        'filename': caller_frame[1],
+                        'line':     caller_frame[2],
+                        }
 
             return f
         return decorator

tests/test_autodoc.py

@@ -1,3 +1,5 @@
+import inspect
+import os.path
 import unittest
 import sys
 import os
@@ -13,6 +15,14 @@ def setUp(self):
         self.app.debug = True
         self.autodoc = Autodoc(self.app)
 
+    @staticmethod
+    def thisFile():
+        """Returns the basename of __file__ without a trailing 'c'"""
+        filename = os.path.basename(__file__)
+        if filename.endswith('.pyc'):
+            filename = filename[:-1]
+        return filename
+
     def testGet(self):
         @self.app.route('/')
         @self.autodoc.doc()
@@ -30,7 +40,7 @@ def index():
         self.assertEqual(d['endpoint'], 'index')
         self.assertEqual(d['docstring'], 'Returns a hello world message')
         self.assertIsInstance(d['location']['line'], int)
-        self.assertIn(__file__, d['location']['filename'])
+        self.assertIn(self.thisFile(), d['location']['filename'])
         self.assertFalse(d['defaults'])
 
     def testPost(self):
@@ -50,7 +60,7 @@ def index():
             self.assertEqual(d['endpoint'], 'index')
             self.assertEqual(d['docstring'], 'Returns a hello world message')
             self.assertIsInstance(d['location']['line'], int)
-            self.assertIn(__file__, d['location']['filename'])
+            self.assertIn(self.thisFile(), d['location']['filename'])
             self.assertFalse(d['defaults'])
 
     def testParams(self):
@@ -248,3 +258,44 @@ def ab(param1, param2):
                 )
             self.assertIn('Returns arguments', doc)
 
+    def testLocation(self):
+        line_no = inspect.stack()[0][2] + 2 # the doc() line
+        @self.app.route('/location')
+        @self.autodoc.doc()
+        def location():
+            return 'location'
+
+        with self.app.app_context():
+            doc = self.autodoc.generate()
+            d = doc[0]
+            self.assertIsInstance(d['location']['line'], int)
+            self.assertEqual(d['location']['line'], line_no)
+            self.assertIn(self.thisFile(), d['location']['filename'])
+
+    def testNoLocation(self):
+        @self.app.route('/location')
+        @self.autodoc.doc(set_location=False)
+        def location():
+            return 'location'
+
+        with self.app.app_context():
+            doc = self.autodoc.generate()
+            d = doc[0]
+            self.assertIsNone(d['location'])
+
+    def testRedecorate(self):
+        @self.app.route('/redecorate')
+        # add to "all" and "group1"
+        @self.autodoc.doc('group1')
+        def redecorate():
+            return 'redecorate'
+
+        # add to "group2"
+        self.app.view_functions['redecorate'] = self.autodoc.doc('group2')(redecorate)
+
+        with self.app.app_context():
+            self.assertTrue(1 == len(self.autodoc.generate('all')))
+            self.assertTrue(1 == len(self.autodoc.generate('group1')))
+            self.assertTrue(1 == len(self.autodoc.generate('group2')))
+            self.assertFalse(1 == len(self.autodoc.generate('group3')))
+