Skip to content

Commit

Permalink
Implement AntlrOpSelectionVisitor on frontend (#26467)
Browse files Browse the repository at this point in the history
## Summary & Motivation
Implement visitor pattern for op selection syntax.

## How I Tested These Changes
`AntlrOpSelection.test.ts`
  • Loading branch information
briantu authored Dec 13, 2024
1 parent 5d4edc8 commit db1b5f0
Show file tree
Hide file tree
Showing 3 changed files with 283 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
import {CharStreams, CommonTokenStream} from 'antlr4ts';

import {AntlrOpSelectionVisitor} from './AntlrOpSelectionVisitor';
import {GraphQueryItem} from '../app/GraphQueryImpl';
import {AntlrInputErrorListener} from '../asset-selection/AntlrAssetSelection';
import {OpSelectionLexer} from './generated/OpSelectionLexer';
import {OpSelectionParser} from './generated/OpSelectionParser';

type OpSelectionQueryResult = {
all: GraphQueryItem[];
focus: GraphQueryItem[];
};

export const parseOpSelectionQuery = (
all_ops: GraphQueryItem[],
query: string,
): OpSelectionQueryResult | Error => {
try {
const lexer = new OpSelectionLexer(CharStreams.fromString(query));
lexer.removeErrorListeners();
lexer.addErrorListener(new AntlrInputErrorListener());

const tokenStream = new CommonTokenStream(lexer);

const parser = new OpSelectionParser(tokenStream);
parser.removeErrorListeners();
parser.addErrorListener(new AntlrInputErrorListener());

const tree = parser.start();

const visitor = new AntlrOpSelectionVisitor(all_ops);
const all_selection = visitor.visit(tree);
const focus_selection = visitor.focus_ops;

return {
all: Array.from(all_selection),
focus: Array.from(focus_selection),
};
} catch (e) {
return e as Error;
}
};
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
import {AbstractParseTreeVisitor} from 'antlr4ts/tree/AbstractParseTreeVisitor';

import {GraphQueryItem, GraphTraverser} from '../app/GraphQueryImpl';
import {
AllExpressionContext,
AndExpressionContext,
AttributeExpressionContext,
DownTraversalExpressionContext,
NameExprContext,
NameSubstringExprContext,
NotExpressionContext,
OrExpressionContext,
ParenthesizedExpressionContext,
StartContext,
TraversalAllowedExpressionContext,
UpAndDownTraversalExpressionContext,
UpTraversalExpressionContext,
} from './generated/OpSelectionParser';
import {OpSelectionVisitor} from './generated/OpSelectionVisitor';
import {getTraversalDepth, getValue} from '../asset-selection/AntlrAssetSelectionVisitor';

export class AntlrOpSelectionVisitor
extends AbstractParseTreeVisitor<Set<GraphQueryItem>>
implements OpSelectionVisitor<Set<GraphQueryItem>>
{
all_ops: Set<GraphQueryItem>;
focus_ops: Set<GraphQueryItem>;
traverser: GraphTraverser<GraphQueryItem>;

protected defaultResult() {
return new Set<GraphQueryItem>();
}

constructor(all_ops: GraphQueryItem[]) {
super();
this.all_ops = new Set(all_ops);
this.focus_ops = new Set();
this.traverser = new GraphTraverser(all_ops);
}

visitStart(ctx: StartContext) {
return this.visit(ctx.expr());
}

visitTraversalAllowedExpression(ctx: TraversalAllowedExpressionContext) {
return this.visit(ctx.traversalAllowedExpr());
}

visitUpAndDownTraversalExpression(ctx: UpAndDownTraversalExpressionContext) {
const selection = this.visit(ctx.traversalAllowedExpr());
const up_depth: number = getTraversalDepth(ctx.traversal(0));
const down_depth: number = getTraversalDepth(ctx.traversal(1));
const selection_copy = new Set(selection);
for (const item of selection_copy) {
this.traverser.fetchUpstream(item, up_depth).forEach((i) => selection.add(i));
this.traverser.fetchDownstream(item, down_depth).forEach((i) => selection.add(i));
}
return selection;
}

visitUpTraversalExpression(ctx: UpTraversalExpressionContext) {
const selection = this.visit(ctx.traversalAllowedExpr());
const traversal_depth: number = getTraversalDepth(ctx.traversal());
const selection_copy = new Set(selection);
for (const item of selection_copy) {
this.traverser.fetchUpstream(item, traversal_depth).forEach((i) => selection.add(i));
}
return selection;
}

visitDownTraversalExpression(ctx: DownTraversalExpressionContext) {
const selection = this.visit(ctx.traversalAllowedExpr());
const traversal_depth: number = getTraversalDepth(ctx.traversal());
const selection_copy = new Set(selection);
for (const item of selection_copy) {
this.traverser.fetchDownstream(item, traversal_depth).forEach((i) => selection.add(i));
}
return selection;
}

visitNotExpression(ctx: NotExpressionContext) {
const selection = this.visit(ctx.expr());
return new Set([...this.all_ops].filter((i) => !selection.has(i)));
}

visitAndExpression(ctx: AndExpressionContext) {
const left = this.visit(ctx.expr(0));
const right = this.visit(ctx.expr(1));
return new Set([...left].filter((i) => right.has(i)));
}

visitOrExpression(ctx: OrExpressionContext) {
const left = this.visit(ctx.expr(0));
const right = this.visit(ctx.expr(1));
return new Set([...left, ...right]);
}

visitAllExpression(_ctx: AllExpressionContext) {
return this.all_ops;
}

visitAttributeExpression(ctx: AttributeExpressionContext) {
return this.visit(ctx.attributeExpr());
}

visitParenthesizedExpression(ctx: ParenthesizedExpressionContext) {
return this.visit(ctx.expr());
}

visitNameExpr(ctx: NameExprContext) {
const value: string = getValue(ctx.value());
const selection = [...this.all_ops].filter((i) => i.name === value);
selection.forEach((i) => this.focus_ops.add(i));
return new Set(selection);
}

visitNameSubstringExpr(ctx: NameSubstringExprContext) {
const value: string = getValue(ctx.value());
const selection = [...this.all_ops].filter((i) => i.name.includes(value));
selection.forEach((i) => this.focus_ops.add(i));
return new Set(selection);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,118 @@
/* eslint-disable jest/expect-expect */

import {GraphQueryItem} from '../../app/GraphQueryImpl';
import {parseOpSelectionQuery} from '../AntlrOpSelection';

const TEST_GRAPH: GraphQueryItem[] = [
// Top Layer
{
name: 'A',
inputs: [{dependsOn: []}],
outputs: [{dependedBy: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}],
},
// Second Layer
{
name: 'B',
inputs: [{dependsOn: [{solid: {name: 'A'}}]}],
outputs: [{dependedBy: [{solid: {name: 'C'}}]}],
},
{
name: 'B2',
inputs: [{dependsOn: [{solid: {name: 'A'}}]}],
outputs: [{dependedBy: [{solid: {name: 'C'}}]}],
},
// Third Layer
{
name: 'C',
inputs: [{dependsOn: [{solid: {name: 'B'}}, {solid: {name: 'B2'}}]}],
outputs: [{dependedBy: []}],
},
];

function assertQueryResult(query: string, expectedNames: string[]) {
const result = parseOpSelectionQuery(TEST_GRAPH, query);
expect(result).not.toBeInstanceOf(Error);
if (result instanceof Error) {
throw result;
}
expect(result.all.length).toBe(expectedNames.length);
expect(new Set(result.all.map((op) => op.name))).toEqual(new Set(expectedNames));
}

// Most tests copied from AntlrAssetSelection.test.ts
describe('parseOpSelectionQuery', () => {
describe('invalid queries', () => {
it('should throw on invalid queries', () => {
expect(parseOpSelectionQuery(TEST_GRAPH, 'A')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'name:A name:B')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'not')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'and')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'name:A and')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'sinks(*)')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'roots(*)')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'notafunction()')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'tag:foo=')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'owner')).toBeInstanceOf(Error);
expect(parseOpSelectionQuery(TEST_GRAPH, 'owner:[email protected]')).toBeInstanceOf(Error);
});
});

describe('valid queries', () => {
it('should parse star query', () => {
assertQueryResult('*', ['A', 'B', 'B2', 'C']);
});

it('should parse name query', () => {
assertQueryResult('name:A', ['A']);
});

it('should parse name_substring query', () => {
assertQueryResult('name_substring:A', ['A']);
assertQueryResult('name_substring:B', ['B', 'B2']);
});

it('should parse and query', () => {
assertQueryResult('name:A and name:B', []);
assertQueryResult('name:A and name:B and name:C', []);
});

it('should parse or query', () => {
assertQueryResult('name:A or name:B', ['A', 'B']);
assertQueryResult('name:A or name:B or name:C', ['A', 'B', 'C']);
assertQueryResult('(name:A or name:B) and (name:B or name:C)', ['B']);
});

it('should parse upstream plus query', () => {
assertQueryResult('+name:A', ['A']);
assertQueryResult('+name:B', ['A', 'B']);
assertQueryResult('+name:C', ['B', 'B2', 'C']);
assertQueryResult('++name:C', ['A', 'B', 'B2', 'C']);
});

it('should parse downstream plus query', () => {
assertQueryResult('name:A+', ['A', 'B', 'B2']);
assertQueryResult('name:A++', ['A', 'B', 'B2', 'C']);
assertQueryResult('name:C+', ['C']);
assertQueryResult('name:B+', ['B', 'C']);
});

it('should parse upstream star query', () => {
assertQueryResult('*name:A', ['A']);
assertQueryResult('*name:B', ['A', 'B']);
assertQueryResult('*name:C', ['A', 'B', 'B2', 'C']);
});

it('should parse downstream star query', () => {
assertQueryResult('name:A*', ['A', 'B', 'B2', 'C']);
assertQueryResult('name:B*', ['B', 'C']);
assertQueryResult('name:C*', ['C']);
});

it('should parse up and down traversal queries', () => {
assertQueryResult('name:A* and *name:C', ['A', 'B', 'B2', 'C']);
assertQueryResult('*name:B*', ['A', 'B', 'C']);
assertQueryResult('name:A* and *name:C and *name:B*', ['A', 'B', 'C']);
assertQueryResult('name:A* and *name:B* and *name:C', ['A', 'B', 'C']);
});
});
});

1 comment on commit db1b5f0

@github-actions
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deploy preview for dagit-core-storybook ready!

✅ Preview
https://dagit-core-storybook-efu2i3abh-elementl.vercel.app

Built with commit db1b5f0.
This pull request is being automatically deployed with vercel-action

Please sign in to comment.