content: package org.netbeans.modules.python.testrunner;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.regex.MatchResult;
  import java.util.regex.Pattern;
  import java.util.stream.Stream;
  import javax.swing.text.BadLocationException;
  import javax.swing.text.Position;
  import javax.swing.text.StyledDocument;
  import org.netbeans.api.project.FileOwnerQuery;
  import org.netbeans.api.project.Project;
  import org.netbeans.api.project.ProjectUtils;
  import org.netbeans.api.project.SourceGroup;
  import org.netbeans.modules.editor.NbEditorUtilities;
  import org.netbeans.modules.gsf.testrunner.ui.api.TestMethodController.TestMethod;
  import org.netbeans.modules.parsing.api.Snapshot;
  import org.netbeans.spi.project.SingleMethod;
  import org.openide.cookies.EditorCookie;
  import org.openide.filesystems.FileObject;
  import org.openide.filesystems.FileUtil;
  import org.openide.loaders.DataObject;
  import org.openide.loaders.DataObjectNotFoundException;
  import org.openide.text.NbDocument;
  import org.openide.util.Exceptions;
  
  /**
   *
   * @author albilu
   */
  class PythonTestMethodController {
  
      static Pattern TEST_CLASS = Pattern.compile(".*\\nclass\\s+([^\\(|:]+).*:");
      static Pattern TEST_METHOD = Pattern.compile(".*def\\s+(test_[^\\(|:]+).*:");
  
      static void setMethods(Snapshot snapshot) {
          try {
              FileObject fo = snapshot.getSource().getFileObject();
              StyledDocument document = DataObject.find(fo).getLookup().lookup(EditorCookie.class).getDocument();
              List<TestMethod> computeMethods = computeMethods(document);
              if (computeMethods == null || computeMethods.isEmpty()) {
                  return;
              }
              NbDocument.runAtomic(document, () -> {
                  setTestMethodsImpl(document, computeMethods);
              });
  
          } catch (DataObjectNotFoundException ex) {
              Exceptions.printStackTrace(ex);
          }
      }
  
      //Copy from org/netbeans/modules/gsf/testrunner/ui/api/TestMethodController.java
      private static void setTestMethodsImpl(StyledDocument doc, List<TestMethod> methods) {
          //FIXME-BUG Method Run annotation position change break action => contribution-welcome
          doc.putProperty(PythonTestMethodAnnotation.DOCUMENT_METHODS_KEY, methods);
  
          Map<TestMethod, PythonTestMethodAnnotation> annotations = (Map<TestMethod, PythonTestMethodAnnotation>) doc.getProperty(PythonTestMethodAnnotation.DOCUMENT_ANNOTATIONS_KEY);
  
          if (annotations == null) {
              annotations = new HashMap<>();
              doc.putProperty(PythonTestMethodAnnotation.DOCUMENT_ANNOTATIONS_KEY, annotations);
          }
  
          Map<Integer, TestMethod> annotationLines = (Map<Integer, TestMethod>) doc.getProperty(PythonTestMethodAnnotation.DOCUMENT_ANNOTATION_LINES_KEY);
  
          if (annotationLines == null) {
              annotationLines = new HashMap<>();
              doc.putProperty(PythonTestMethodAnnotation.DOCUMENT_ANNOTATION_LINES_KEY, annotationLines);
          }
  
          Map<TestMethod, PythonTestMethodAnnotation> removed = new HashMap<>(annotations);
  
          methods.forEach(tm -> removed.remove(tm));
  
          Set<TestMethod> added = new HashSet<>(methods);
  
          added.removeAll(annotations.keySet());
  
          for (TestMethod method : added) {
              PythonTestMethodAnnotation a = new PythonTestMethodAnnotation(method);
              NbDocument.addAnnotation(doc, method.preferred(), 0, a);
              annotations.put(method, a);
              int line = NbDocument.findLineNumber(doc, method.preferred().getOffset());
              annotationLines.put(line, method);
          }
          for (Map.Entry<TestMethod, PythonTestMethodAnnotation> e : removed.entrySet()) {
              NbDocument.removeAnnotation(doc, e.getValue());
              annotations.remove(e.getKey());
              int line = NbDocument.findLineNumber(doc, e.getKey().preferred().getOffset());
              annotationLines.remove(line);
          }
      }
  
      private static List<TestMethod> computeMethods(StyledDocument doc) {
          if (doc == null) {
              return null;
          }
          FileObject fo = NbEditorUtilities.getFileObject(doc);
          if (!isUnderTestSources(fo)) {
              return null;
          }
          List<TestMethod> methods = new ArrayList<>();
          try {
              String currentFile = doc.getText(0, doc.getEndPosition().getOffset());
              Stream<MatchResult> classResults = TEST_CLASS.matcher(currentFile).results();
  
              if (classResults.count() == 0) {
                  return null;
              }
  
              TEST_METHOD.matcher(currentFile).results().forEach((methodMatch) -> {
                  try {
                      int startOffset = methodMatch.start(1);
                      MatchResult classGroup = TEST_CLASS.matcher(currentFile).results()
                              .filter((classMatch) -> classMatch.end(1) < startOffset)
                              .reduce((first, second) -> second)
                              .orElse(null);
                      if (classGroup != null) {
                          String className = classGroup.group(1);
                          String testMethodName = methodMatch.group(1);
                          Position start = NbDocument.createPosition(doc, startOffset, Position.Bias.Forward);
                          Position end = NbDocument.createPosition(doc, methodMatch.end(1), Position.Bias.Backward);
                          methods.add(new TestMethod(className, new SingleMethod(fo, testMethodName), start, start, end));
                      }
                  } catch (BadLocationException ex) {
                      Exceptions.printStackTrace(ex);
                  }
  
              });
          } catch (BadLocationException ex) {
              Exceptions.printStackTrace(ex);
          }
  
          return methods;
      }
  
      private static boolean isUnderTestSources(FileObject fo) {
          Project owner = FileOwnerQuery.getOwner(fo);
          if (owner == null) {
              return false;
          }
  
          for (SourceGroup dir : ProjectUtils.getSources(owner).getSourceGroups("testsources")) {
              if (FileUtil.isParentOf(dir.getRootFolder(), fo)) {
                  return true;
              }
          }
          return false;
      }
  }
  