feat: add support for Prefect 3.0 pipelines#26
Open
Resmung0 wants to merge 2 commits into
Open
Conversation
- Implement PrefectParser to detect and parse Prefect flows and tasks from Python files. - Support task dependencies via data passing and `wait_for` arguments. - Register PrefectParser in DataProcessingPipeline. - Add unit tests for PrefectParser. - Add ts-jest to devDependencies to support TypeScript tests. Co-authored-by: Resmung0 <36093882+Resmung0@users.noreply.github.com>
Contributor
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
Reviewer's GuideAdds a new Prefect 3.0 data-processing parser and integrates it into the existing pipeline infrastructure, including unit tests and a minor test tooling dependency bump. Sequence diagram for PrefectParser handling a Prefect 3.0 Python filesequenceDiagram
participant VSCodeExtension
participant DataProcessingPipeline
participant PrefectParser
participant VisualizationEngine
VSCodeExtension->>DataProcessingPipeline: analyzeWorkspace()
loop for each file
DataProcessingPipeline->>PrefectParser: canParse(fileName, content)
alt file is Prefect flow
PrefectParser-->>DataProcessingPipeline: true
DataProcessingPipeline->>PrefectParser: parse(content, filePath)
PrefectParser->>PrefectParser: extract taskNames via taskRegex
PrefectParser->>PrefectParser: extract flowBody via flowRegex
PrefectParser->>PrefectParser: build nodes and edges
PrefectParser-->>DataProcessingPipeline: PipelineData
DataProcessingPipeline->>VisualizationEngine: registerPipeline(PipelineData)
else not a Prefect file
PrefectParser-->>DataProcessingPipeline: false
end
end
Class diagram for PrefectParser integration into DataProcessingPipelineclassDiagram
class DataProcessingPipeline {
+PipelinePatternType type
+IParser[] parsers
}
class IParser {
<<interface>>
+string name
+canParse(fileName string, content string) boolean
+parse(content string, filePath string) Promise~PipelineData~
}
class PrefectParser {
+string name
+canParse(fileName string, content string) boolean
+parse(content string, filePath string) Promise~PipelineData~
}
class AirflowParser {
+string name
+canParse(fileName string, content string) boolean
+parse(content string, filePath string) Promise~PipelineData~
}
class KedroParser {
+string name
+canParse(fileName string, content string) boolean
+parse(content string, filePath string) Promise~PipelineData~
}
class DVCParser {
+string name
+canParse(fileName string, content string) boolean
+parse(content string, filePath string) Promise~PipelineData~
}
class PipelineData {
+string filePath
+string framework
+PipelineNode[] nodes
+PipelineEdge[] edges
}
class PipelineNode {
+string id
+string label
+string type
+any data
}
class PipelineEdge {
+string id
+string source
+string target
}
DataProcessingPipeline o-- IParser : uses
DataProcessingPipeline o-- PrefectParser : parser
DataProcessingPipeline o-- AirflowParser : parser
DataProcessingPipeline o-- KedroParser : parser
DataProcessingPipeline o-- DVCParser : parser
PrefectParser ..|> IParser
AirflowParser ..|> IParser
KedroParser ..|> IParser
DVCParser ..|> IParser
PrefectParser --> PipelineData
PipelineData --> PipelineNode
PipelineData --> PipelineEdge
Flow diagram for PrefectParser parse logic and dependency detectionflowchart TD
A["Start parse(content, filePath)"] --> B["Initialize empty nodes and edges"]
B --> C["Find @task decorated functions with taskRegex"]
C --> D["For each taskName: add PipelineNode and store in taskNames"]
D --> E["Find first @flow definition and extract flowBody with flowRegex"]
E --> F{Flow found?}
F -- "No" --> G["Return PipelineData with only task nodes"]
F -- "Yes" --> H["Initialize varToTask map"]
H --> I["Split flowBody into lines"]
I --> J["Iterate lines and taskNames"]
J --> K["Match assignments: var = taskName.submit(...) or var = taskName(...)"]
K --> L["On match: map variable to taskName in varToTask"]
L --> M["Match calls: taskName.submit(args) or taskName(args)"]
M --> N{callMatch exists?}
N -- "No" --> J
N -- "Yes" --> O["Inspect args for variables present in varToTask"]
O --> P["For each matching variable: add edge sourceTask -> taskName if not exists"]
P --> Q["Inspect args for nested calls otherTask(')"]
Q --> R["For each otherTask in taskNames: if args contains otherTask( add edge otherTask -> taskName"]
R --> J
G --> S["Return PipelineData(filePath, framework, nodes, edges)"]
J --> T["All lines processed"]
T --> S
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The
PrefectParsercurrently only processes the first@flowmatch in a file (flowRegex.execused once), so additional flows in the same module will be ignored; consider iterating over all matches to capture multiple flows. - The
canParseheuristic (content.includes("from prefect import") && content.includes("@flow")) will miss valid Prefect 3 flows that use other import styles (e.g.,import prefect,from prefect import flowwithouttask, orfrom prefect import flow as f); consider loosening this detection to rely on broaderprefectimports and/or decorator patterns.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `PrefectParser` currently only processes the first `@flow` match in a file (`flowRegex.exec` used once), so additional flows in the same module will be ignored; consider iterating over all matches to capture multiple flows.
- The `canParse` heuristic (`content.includes("from prefect import") && content.includes("@flow")`) will miss valid Prefect 3 flows that use other import styles (e.g., `import prefect`, `from prefect import flow` without `task`, or `from prefect import flow as f`); consider loosening this detection to rely on broader `prefect` imports and/or decorator patterns.
## Individual Comments
### Comment 1
<location path="src/extension/parsers/data-processing/PrefectParser.ts" line_range="35-36" />
<code_context>
+ }
+
+ // 2. Identify the flow and its body
+ const flowRegex = /@flow(?:\([^)]*\))?\s+def\s+(\w+)\s*\([^)]*\):([\s\S]+?)(?=\n\S|$)/g;
+ const flowMatch = flowRegex.exec(content);
+ if (flowMatch) {
+ const flowBody = flowMatch[2];
</code_context>
<issue_to_address>
**issue (bug_risk):** Only the first flow definition is parsed even though the regex is global.
`flowRegex` uses the `g` flag but is only executed once, so only the first `@flow` is handled and later flows are skipped. If multiple flows per file are expected, loop over the matches (e.g. `while ((flowMatch = flowRegex.exec(content)) !== null) { ... }`) and aggregate nodes/edges for each flow.
</issue_to_address>
### Comment 2
<location path="src/extension/parsers/data-processing/PrefectParser.ts" line_range="61-64" />
<code_context>
+ const callRegex = new RegExp(`${taskName}(?:\\.submit)?\\s*\\(([^)]*)\\)`);
+ const callMatch = line.match(callRegex);
+ if (callMatch) {
+ const args = callMatch[1];
+
+ // Check for variables in args that map to tasks
+ varToTask.forEach((sourceTask, variable) => {
+ if (args.includes(variable)) {
+ const edgeId = `e-${sourceTask}-${taskName}`;
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Variable-to-task matching via `includes` can cause false-positive edges.
`args.includes(variable)` risks matching substrings or content inside string literals (e.g., variable `a` matching `data`), which can produce spurious dependencies. It would be safer to parse the argument list (e.g., splitting into identifiers) or at least match with word boundaries via a regex (e.g., `new RegExp(`\b${variable}\b`)`).
</issue_to_address>
### Comment 3
<location path="tests/PrefectParser.test.ts" line_range="17" />
<code_context>
+ expect(parser.canParse('flow.txt', content)).toBe(false);
+ });
+
+ it('should parse tasks and dependencies from a flow', async () => {
+ const content = `
+from prefect import flow, task
</code_context>
<issue_to_address>
**suggestion (testing):** Add a test case for dependencies created via plain task calls (without `.submit()`), since the parser supports both.
The assignment regex already supports both `task_name.submit(...)` and `task_name(...)`, but the current dependency test only exercises `a = task_a.submit()`. Please add coverage for the plain-call style—e.g., `a = task_a()` and `b = task_b(a)`—and assert the same edge structure, so we don’t regress support for non-`.submit()` calls.
Suggested implementation:
```typescript
it('should parse tasks and dependencies from a flow', async () => {
const content = `
from prefect import flow, task
@task
def task_a():
return "a"
@task
def task_b(data):
return "b"
@task
def task_c():
return "c"
@flow
def submit_style_flow():
a = task_a.submit()
b = task_b.submit(a)
c = task_c.submit()
return b
@flow
def direct_call_flow():
a = task_a()
b = task_b(a)
c = task_c()
return b
```
The updated flow content now includes **both** dependency styles:
- `submit_style_flow` uses `task_a.submit(...)` / `task_b.submit(a)`.
- `direct_call_flow` uses plain calls `task_a(...)` / `task_b(a)`.
To fully implement the test intent, you should verify that the parser identifies edges for *both* styles. Depending on how `PrefectParser.parse` currently reports tasks/edges, you’ll likely need to:
1. Ensure the expectations in this test assert that there is an edge from `task_a` to `task_b` created from:
- The `.submit()` call, and
- The direct call.
2. If the parser groups all edges together without distinguishing by flow function, the test can simply assert that there is at least one `task_a -> task_b` edge for each style (or that there are two such edges, if the parser does not deduplicate).
Adjust the expectations in this test accordingly using your existing assertion pattern for dependencies. If you’d prefer truly separate test cases (one for `.submit` and one for direct calls), you can split this combined flow into two distinct tests, each using only one of the two flow definitions above, and reuse the same expectations in both tests.
</issue_to_address>
### Comment 4
<location path="src/extension/parsers/data-processing/PrefectParser.ts" line_range="15" />
<code_context>
+ );
+ }
+
+ async parse(content: string, filePath: string): Promise<PipelineData> {
+ const nodes: PipelineNode[] = [];
+ const edges: PipelineEdge[] = [];
</code_context>
<issue_to_address>
**issue (complexity):** Consider refactoring the parse logic into small helper methods with separate passes for collecting assignments and building edges while reusing precomputed regexes instead of creating them per line.
You can keep the behavior but reduce complexity by:
- splitting the `parse` logic into small helpers, and
- separating “assignment collection” from “edge extraction,”
- avoiding per-line, per-task regex construction.
### 1. Extract helpers for clarity
Move task and flow extraction into dedicated functions so `parse` reads as a high-level pipeline:
```ts
async parse(content: string, filePath: string): Promise<PipelineData> {
const nodes: PipelineNode[] = [];
const edges: PipelineEdge[] = [];
const taskNames = this.extractTasks(content, nodes);
const flowBody = this.extractFlowBody(content);
if (flowBody) {
const lines = flowBody.split('\n');
const varToTask = this.buildVarToTaskMap(lines, taskNames);
this.extractEdgesFromLines(lines, taskNames, varToTask, edges);
}
return { filePath, framework: this.name, nodes, edges };
}
private extractTasks(content: string, nodes: PipelineNode[]): string[] {
const taskRegex = /@task(?:\([^)]*\))?\s+def\s+(\w+)/g;
const taskNames: string[] = [];
let match;
while ((match = taskRegex.exec(content)) !== null) {
const name = match[1];
taskNames.push(name);
nodes.push({ id: name, label: name, type: 'default', data: { framework: this.name } });
}
return taskNames;
}
private extractFlowBody(content: string): string | null {
const flowRegex = /@flow(?:\([^)]*\))?\s+def\s+(\w+)\s*\([^)]*\):([\s\S]+?)(?=\n\S|$)/g;
const flowMatch = flowRegex.exec(content);
return flowMatch ? flowMatch[2] : null;
}
```
### 2. Two passes: assignments vs edges
This reduces interleaving state/side‑effects and removes one level of nesting:
```ts
private buildVarToTaskMap(lines: string[], taskNames: string[]): Map<string, string> {
const varToTask = new Map<string, string>();
// Precompute assignment regexes per task
const assignmentRegexes = taskNames.map(
t => ({ task: t, re: new RegExp(`(\\w+)\\s*=\\s*${t}(?:\\.submit)?\\s*\\(`) })
);
for (const line of lines) {
for (const { task, re } of assignmentRegexes) {
const m = line.match(re);
if (m) {
varToTask.set(m[1], task);
}
}
}
return varToTask;
}
private extractEdgesFromLines(
lines: string[],
taskNames: string[],
varToTask: Map<string, string>,
edges: PipelineEdge[]
): void {
// Precompute call regexes per task
const callRegexes = taskNames.map(
t => ({ task: t, re: new RegExp(`${t}(?:\\.submit)?\\s*\\(([^)]*)\\)`) })
);
const ensureEdge = (source: string, target: string) => {
const id = `e-${source}-${target}`;
if (!edges.some(e => e.id === id)) {
edges.push({ id, source, target });
}
};
for (const line of lines) {
for (const { task: targetTask, re } of callRegexes) {
const callMatch = line.match(re);
if (!callMatch) continue;
const args = callMatch[1];
// vars -> tasks
varToTask.forEach((sourceTask, variable) => {
if (args.includes(variable)) {
ensureEdge(sourceTask, targetTask);
}
});
// direct nested calls: task2(task1())
for (const otherTask of taskNames) {
if (otherTask !== targetTask && args.includes(`${otherTask}(`)) {
ensureEdge(otherTask, targetTask);
}
}
}
}
}
```
This preserves the current behavior but:
- removes per‑line regex allocation,
- isolates responsibilities into small, testable helpers,
- and simplifies control flow (no mutation of `varToTask` while resolving dependencies).
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| expect(parser.canParse('flow.txt', content)).toBe(false); | ||
| }); | ||
|
|
||
| it('should parse tasks and dependencies from a flow', async () => { |
There was a problem hiding this comment.
suggestion (testing): Add a test case for dependencies created via plain task calls (without .submit()), since the parser supports both.
The assignment regex already supports both task_name.submit(...) and task_name(...), but the current dependency test only exercises a = task_a.submit(). Please add coverage for the plain-call style—e.g., a = task_a() and b = task_b(a)—and assert the same edge structure, so we don’t regress support for non-.submit() calls.
Suggested implementation:
it('should parse tasks and dependencies from a flow', async () => {
const content = `
from prefect import flow, task
@task
def task_a():
return "a"
@task
def task_b(data):
return "b"
@task
def task_c():
return "c"
@flow
def submit_style_flow():
a = task_a.submit()
b = task_b.submit(a)
c = task_c.submit()
return b
@flow
def direct_call_flow():
a = task_a()
b = task_b(a)
c = task_c()
return b
The updated flow content now includes both dependency styles:
submit_style_flowusestask_a.submit(...)/task_b.submit(a).direct_call_flowuses plain callstask_a(...)/task_b(a).
To fully implement the test intent, you should verify that the parser identifies edges for both styles. Depending on how PrefectParser.parse currently reports tasks/edges, you’ll likely need to:
- Ensure the expectations in this test assert that there is an edge from
task_atotask_bcreated from:- The
.submit()call, and - The direct call.
- The
- If the parser groups all edges together without distinguishing by flow function, the test can simply assert that there is at least one
task_a -> task_bedge for each style (or that there are two such edges, if the parser does not deduplicate).
Adjust the expectations in this test accordingly using your existing assertion pattern for dependencies. If you’d prefer truly separate test cases (one for .submit and one for direct calls), you can split this combined flow into two distinct tests, each using only one of the two flow definitions above, and reuse the same expectations in both tests.
- Support multiple flows per file. - Improve 'canParse' to handle various import and decorator styles (e.g., @prefect.flow, @f). - Refactor 'parse' into helper methods for better maintainability and performance. - Use word boundaries in variable matching to avoid false-positive dependencies. - Optimize regex usage by precomputing them. - Add comprehensive unit tests for new features and edge cases. Co-authored-by: Resmung0 <36093882+Resmung0@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds support for Prefect 3.0 pipelines to the Caldera extension.
Key changes:
PrefectParserinsrc/extension/parsers/data-processing/which uses static analysis (regex) to identify@flowand@taskdecorated functions in Python files..submit()calls) and checkingwait_forparameters.PrefectParserin theDataProcessingPipelinefor**/*.pyfiles.tests/PrefectParser.test.tsto verify flow detection and dependency resolution.package.jsonto includets-jestfor running TypeScript-based unit tests.The implementation follows the established pattern for data processing parsers and integrates with the existing visualization infrastructure.
PR created automatically by Jules for task 15643003390712396561 started by @Resmung0
Summary by Sourcery
Add a Prefect data-processing parser and register it in the pipeline to visualize Prefect-based Python workflows, alongside updating testing dependencies.
New Features:
Enhancements:
Build: