forked from DetachHead/rebased
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathPyCellUtil.java
More file actions
125 lines (106 loc) · 3.94 KB
/
Copy pathPyCellUtil.java
File metadata and controls
125 lines (106 loc) · 3.94 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.jetbrains.python;
import com.intellij.openapi.util.TextRange;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.psi.PsiComment;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.util.PsiTreeUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
public final class PyCellUtil {
public static final class CellRange {
public final TextRange textRange;
public final Boolean lastCell;
public CellRange(TextRange textRange, Boolean lastCell) {
this.textRange = textRange;
this.lastCell = lastCell;
}
}
public static boolean hasCells(@NotNull PsiFile file) {
PsiElement el = file.getFirstChild();
while (el != null) {
if (isBlockCell(el)) {
return true;
}
el = el.getNextSibling();
}
return false;
}
public static @Nullable PsiElement getCellStart(@NotNull PsiElement element) {
PsiElement el = element;
while (el != null && !isBlockCell(el)) {
el = PsiTreeUtil.prevLeaf(el);
}
if (el == null) {
return element.getContainingFile().getFirstChild();
}
return PsiTreeUtil.nextLeaf(el);
}
public static final String[] CELL_SEPARATORS = {
"# %%",
"#%%",
"# <codecell>",
"# In[",
"# COMMAND ----------" // Databricks py-notebooks separator.
};
public static boolean isBlockDefinition(String text) {
for (String separator : CELL_SEPARATORS) {
if (text.startsWith(separator)) {
return true;
}
}
return false;
}
public static boolean isBlockCell(PsiElement element) {
return (element instanceof PsiComment) && isBlockDefinition(element.getText());
}
public static @NotNull String getCodeInCell(@Nullable PsiElement element) {
StringBuilder text = new StringBuilder();
while (element != null && !isBlockCell(element)) {
text.append(element.getText());
element = element.getNextSibling();
}
return StringUtil.trim(text.toString());
}
public static CellRange getCellTextRangeIncludingSeparators(@NotNull PsiElement element) {
PsiElement el = element;
while (el != null && !isBlockCell(el)) {
el = PsiTreeUtil.prevLeaf(el);
}
int startOffset = el == null ? element.getContainingFile().getFirstChild().getTextOffset() : el.getTextRange().getStartOffset();
PsiElement nextCell = findNextCell(element);
int endOffset = nextCell == null ? element.getContainingFile().getLastChild().getTextRange().getEndOffset()
: nextCell.getTextRange().getStartOffset();
return new CellRange(new TextRange(startOffset, endOffset), nextCell == null);
}
public static @Nullable PsiElement findNextCell(PsiElement startElement) {
PsiElement el = PsiTreeUtil.nextLeaf(PsiTreeUtil.getDeepestFirst(startElement));
while (el != null) {
if (isBlockCell(el)) {
return el;
}
el = PsiTreeUtil.nextLeaf(el);
}
return null;
}
public static @Nullable PsiElement findSeparatorOfPrevCell(PsiElement startElement) {
PsiElement currentCellSeparator = PsiTreeUtil.prevLeaf(PsiTreeUtil.getDeepestFirst(startElement));
while (currentCellSeparator != null && !isBlockCell(currentCellSeparator)) {
currentCellSeparator = PsiTreeUtil.prevLeaf(currentCellSeparator);
}
if (currentCellSeparator == null) return null;
return findPrevCell(currentCellSeparator);
}
// Not a prev cell, but only a cell separator that is over the startElement (it could be in the same cell).
public static @Nullable PsiElement findPrevCell(PsiElement startElement) {
PsiElement el = PsiTreeUtil.prevLeaf(PsiTreeUtil.getDeepestFirst(startElement));
while (el != null) {
if (isBlockCell(el)) {
return el;
}
el = PsiTreeUtil.prevLeaf(el);
}
return null;
}
}