18423 source view

netbox/extras/scripts.py

@@ -2,10 +2,12 @@ import inspect
 import json
 import logging
 import os
+import re
 
 import yaml
 from django import forms
 from django.conf import settings
+from django.core.files.storage import storages
 from django.core.validators import RegexValidator
 from django.utils import timezone
 from django.utils.functional import classproperty
@@ -367,9 +369,43 @@ class BaseScript:
     def filename(self):
         return inspect.getfile(self.__class__)
 
+    def findsource(self, object):
+        storage = storages.create_storage(storages.backends["scripts"])
+        with storage.open(os.path.basename(self.filename), 'r') as f:
+            data = f.read()
+
+        # Break the source code into lines
+        lines = [line + '\n' for line in data.splitlines()]
+
+        # Find the class definition
+        name = object.__name__
+        pat = re.compile(r'^(\s*)class\s*' + name + r'\b')
+        # use the class definition with the least indentation
+        candidates = []
+        for i in range(len(lines)):
+            match = pat.match(lines[i])
+            if match:
+                if lines[i][0] == 'c':
+                    return lines, i
+
+                candidates.append((match.group(1), i))
+        if not candidates:
+            raise OSError('could not find class definition')
+
+        # Sort the candidates by whitespace, and by line number
+        candidates.sort()
+        return lines, candidates[0][1]
+
     @property
     def source(self):
-        return inspect.getsource(self.__class__)
+        # Can't use inspect.getsource() as it uses os to get the file
+        # inspect uses ast, but that is overkill for this as we only do
+        # classes.
+        object = self.__class__
+
+        lines, lnum = self.findsource(object)
+        lines = inspect.getblock(lines[lnum:])
+        return ''.join(lines)
 
     @classmethod
     def _get_vars(cls):