@@ -1034,9 +1034,13 @@ class ClassFoundException(Exception):
1034
1034
1035
1035
class _ClassFinder (ast .NodeVisitor ):
1036
1036
1037
- def __init__ (self , qualname ):
1037
+ def __init__ (self , cls , tree , lines , qualname ):
1038
1038
self .stack = []
1039
+ self .cls = cls
1040
+ self .tree = tree
1041
+ self .lines = lines
1039
1042
self .qualname = qualname
1043
+ self .lineno_found = []
1040
1044
1041
1045
def visit_FunctionDef (self , node ):
1042
1046
self .stack .append (node .name )
@@ -1057,11 +1061,48 @@ def visit_ClassDef(self, node):
1057
1061
line_number = node .lineno
1058
1062
1059
1063
# decrement by one since lines starts with indexing by zero
1060
- line_number -= 1
1061
- raise ClassFoundException (line_number )
1064
+ self .lineno_found .append ((line_number - 1 , node .end_lineno ))
1062
1065
self .generic_visit (node )
1063
1066
self .stack .pop ()
1064
1067
1068
+ def get_lineno (self ):
1069
+ self .visit (self .tree )
1070
+ lineno_found_number = len (self .lineno_found )
1071
+ if lineno_found_number == 0 :
1072
+ raise OSError ('could not find class definition' )
1073
+ elif lineno_found_number == 1 :
1074
+ return self .lineno_found [0 ][0 ]
1075
+ else :
1076
+ # We have multiple candidates for the class definition.
1077
+ # Now we have to guess.
1078
+
1079
+ # First, let's see if there are any method definitions
1080
+ for member in self .cls .__dict__ .values ():
1081
+ if isinstance (member , types .FunctionType ):
1082
+ for lineno , end_lineno in self .lineno_found :
1083
+ if lineno <= member .__code__ .co_firstlineno <= end_lineno :
1084
+ return lineno
1085
+
1086
+ class_strings = [('' .join (self .lines [lineno : end_lineno ]), lineno )
1087
+ for lineno , end_lineno in self .lineno_found ]
1088
+
1089
+ # Maybe the class has a docstring and it's unique?
1090
+ if self .cls .__doc__ :
1091
+ ret = None
1092
+ for candidate , lineno in class_strings :
1093
+ if self .cls .__doc__ .strip () in candidate :
1094
+ if ret is None :
1095
+ ret = lineno
1096
+ else :
1097
+ break
1098
+ else :
1099
+ if ret is not None :
1100
+ return ret
1101
+
1102
+ # We are out of ideas, just return the last one found, which is
1103
+ # slightly better than previous ones
1104
+ return self .lineno_found [- 1 ][0 ]
1105
+
1065
1106
1066
1107
def findsource (object ):
1067
1108
"""Return the entire source file and starting line number for an object.
@@ -1098,14 +1139,8 @@ def findsource(object):
1098
1139
qualname = object .__qualname__
1099
1140
source = '' .join (lines )
1100
1141
tree = ast .parse (source )
1101
- class_finder = _ClassFinder (qualname )
1102
- try :
1103
- class_finder .visit (tree )
1104
- except ClassFoundException as e :
1105
- line_number = e .args [0 ]
1106
- return lines , line_number
1107
- else :
1108
- raise OSError ('could not find class definition' )
1142
+ class_finder = _ClassFinder (object , tree , lines , qualname )
1143
+ return lines , class_finder .get_lineno ()
1109
1144
1110
1145
if ismethod (object ):
1111
1146
object = object .__func__
0 commit comments