Skip to content

Commit

Permalink
Add DvcViewer class
Browse files Browse the repository at this point in the history
  • Loading branch information
mattseddon committed Feb 27, 2023
1 parent 9149c01 commit 616c35e
Show file tree
Hide file tree
Showing 8 changed files with 193 additions and 13 deletions.
2 changes: 2 additions & 0 deletions extension/src/cli/dvc/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ export enum SubCommand {

export enum Flag {
ALL_COMMITS = '-A',
FOLLOW = '-f',
FORCE = '-f',
GRANULAR = '--granular',
JOBS = '-j',
Expand All @@ -67,6 +68,7 @@ export enum ExperimentSubCommand {

export enum QueueSubCommand {
KILL = 'kill',
LOGS = 'logs',
START = 'start',
STOP = 'stop'
}
Expand Down
6 changes: 5 additions & 1 deletion extension/src/cli/dvc/runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,11 @@ export class DvcRunner extends Disposable implements ICli {
)

this.pseudoTerminal = this.dispose.track(
new PseudoTerminal(this.processOutput, this.processTerminated)
new PseudoTerminal(
this.processOutput,
this.processTerminated,
'DVC: exp run'
)
)
}

Expand Down
114 changes: 114 additions & 0 deletions extension/src/cli/dvc/viewer.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
import { EventEmitter, Event } from 'vscode'
import { Args, Command, Flag, QueueSubCommand } from './constants'
import { getOptions } from './options'
import { CliResult, CliStarted, ICli, typeCheckCommands } from '..'
import { Config } from '../../config'
import { Disposable } from '../../class/dispose'
import { ViewableCliProcess } from '../viewable'

export const autoRegisteredCommands = {
QUEUE_LOGS: 'queueLogs'
} as const

export class DvcViewer extends Disposable implements ICli {
public readonly autoRegisteredCommands = typeCheckCommands(
autoRegisteredCommands,
this
)

public readonly processCompleted: EventEmitter<CliResult>
public readonly onDidCompleteProcess: Event<CliResult>

public readonly processStarted: EventEmitter<CliStarted>
public readonly onDidStartProcess: Event<CliStarted>

private processes: {
[id: string]: ViewableCliProcess
}

private readonly config: Config

constructor(config: Config) {
super()

this.config = config

this.processes = {}

this.processCompleted = this.dispose.track(new EventEmitter<CliResult>())
this.onDidCompleteProcess = this.processCompleted.event

this.processStarted = this.dispose.track(new EventEmitter<CliStarted>())
this.onDidStartProcess = this.processStarted.event
}

public run(name: string, cwd: string, ...args: Args) {
const viewableProcess = this.getRunningProcess(cwd, ...args)
if (viewableProcess) {
return viewableProcess.show()
}

return this.createProcess(name, cwd, args)
}

public queueLogs(cwd: string, expName: string) {
return this.run(
expName,
cwd,
Command.QUEUE,
QueueSubCommand.LOGS,
expName,
Flag.FOLLOW
)
}

private createProcess(name: string, cwd: string, args: Args) {
const viewableProcess = this.viewProcess(name, cwd, args)

this.setRunningProcess(viewableProcess, cwd, ...args)

const listener = this.dispose.track(
viewableProcess.onDidDispose(() => {
delete this.processes[this.getId(cwd, ...args)]
this.dispose.untrack(listener)
listener.dispose()
})
)
}

private viewProcess(name: string, cwd: string, args: Args) {
return this.dispose.track(
new ViewableCliProcess(
`DVC: ${name}`,
this.getOptions(cwd, args),
this.processStarted,
this.processCompleted
)
)
}

private getOptions(cwd: string, args: Args) {
return getOptions(
this.config.getPythonBinPath(),
this.config.getCliPath(),
cwd,
...args
)
}

private getRunningProcess(cwd: string, ...args: Args) {
return this.processes[this.getId(cwd, ...args)]
}

private setRunningProcess(
viewableProcess: ViewableCliProcess,
cwd: string,
...args: Args
) {
this.processes[this.getId(cwd, ...args)] = viewableProcess
}

private getId(cwd: string, ...args: Args) {
return [cwd, ...args].join(':')
}
}
4 changes: 2 additions & 2 deletions extension/src/cli/viewable.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ export class ViewableCliProcess extends DeferredDisposable {
private readonly pseudoTerminal: PseudoTerminal

constructor(
termName: string,
id: string,
options: ProcessOptions,
processStarted: EventEmitter<CliStarted>,
processCompleted: EventEmitter<CliResult>
Expand All @@ -29,7 +29,7 @@ export class ViewableCliProcess extends DeferredDisposable {
const onDidCloseTerminal = terminalClosed.event

this.pseudoTerminal = this.dispose.track(
new PseudoTerminal(processOutput, terminalClosed, termName)
new PseudoTerminal(processOutput, terminalClosed, id)
)

this.pseudoTerminal.setBlocked(true)
Expand Down
21 changes: 12 additions & 9 deletions extension/src/commands/internal.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@ import { commands } from 'vscode'
import { RegisteredCliCommands, RegisteredCommands } from './external'
import { ICli } from '../cli'
import { Args } from '../cli/constants'
import { autoRegisteredCommands as CliExecutorCommands } from '../cli/dvc/executor'
import { autoRegisteredCommands as CliReaderCommands } from '../cli/dvc/reader'
import { autoRegisteredCommands as dvcRunnerCommands } from '../cli/dvc/runner'
import { autoRegisteredCommands as DvcExecutorCommands } from '../cli/dvc/executor'
import { autoRegisteredCommands as DvcReaderCommands } from '../cli/dvc/reader'
import { autoRegisteredCommands as DvcRunnerCommands } from '../cli/dvc/runner'
import { autoRegisteredCommands as DvcViewerCommands } from '../cli/dvc/viewer'
import { autoRegisteredCommands as GitExecutorCommands } from '../cli/git/executor'
import { autoRegisteredCommands as GitReaderCommands } from '../cli/git/reader'
import { sendTelemetryEvent, sendTelemetryEventAndThrow } from '../telemetry'
Expand All @@ -18,14 +19,16 @@ type Command = (...args: Args) => unknown | Promise<unknown>

export const AvailableCommands = Object.assign(
{ EXP_PUSH: 'expPush' } as const,
CliExecutorCommands,
CliReaderCommands,
dvcRunnerCommands,
DvcExecutorCommands,
DvcReaderCommands,
DvcRunnerCommands,
DvcViewerCommands,
GitExecutorCommands,
GitReaderCommands
) as typeof CliExecutorCommands &
typeof CliReaderCommands &
typeof dvcRunnerCommands &
) as typeof DvcExecutorCommands &
typeof DvcReaderCommands &
typeof DvcRunnerCommands &
typeof DvcViewerCommands &
typeof GitExecutorCommands &
typeof GitReaderCommands & { EXP_PUSH: 'expPush' }
export type CommandId =
Expand Down
4 changes: 4 additions & 0 deletions extension/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ import { Flag } from './cli/dvc/constants'
import { LanguageClient } from './languageClient'
import { collectRunningExperimentPids } from './experiments/processExecution/collect'
import { registerPatchCommand } from './patch'
import { DvcViewer } from './cli/dvc/viewer'
export class Extension extends Disposable {
protected readonly internalCommands: InternalCommands

Expand All @@ -68,6 +69,7 @@ export class Extension extends Disposable {
private readonly dvcExecutor: DvcExecutor
private readonly dvcReader: DvcReader
private readonly dvcRunner: DvcRunner
private readonly dvcViewer: DvcViewer
private readonly gitExecutor: GitExecutor
private readonly gitReader: GitReader

Expand Down Expand Up @@ -95,6 +97,7 @@ export class Extension extends Disposable {
this.dvcExecutor = this.dispose.track(new DvcExecutor(config))
this.dvcReader = this.dispose.track(new DvcReader(config))
this.dvcRunner = this.dispose.track(new DvcRunner(config))
this.dvcViewer = this.dispose.track(new DvcViewer(config))

this.gitExecutor = this.dispose.track(new GitExecutor())
this.gitReader = this.dispose.track(new GitReader())
Expand All @@ -103,6 +106,7 @@ export class Extension extends Disposable {
this.dvcExecutor,
this.dvcReader,
this.dvcRunner,
this.dvcViewer,
this.gitExecutor,
this.gitReader
]
Expand Down
53 changes: 53 additions & 0 deletions extension/src/test/suite/cli/dvc/viewer.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
import { afterEach, beforeEach, describe, it, suite } from 'mocha'
import { expect } from 'chai'
import { restore, spy } from 'sinon'
import { Disposable } from '../../../../extension'
import { Config } from '../../../../config'
import { DvcViewer } from '../../../../cli/dvc/viewer'
import { dvcDemoPath } from '../../../util'
import { ViewableCliProcess } from '../../../../cli/viewable'

suite('DVC Viewer Test Suite', () => {
const disposable = Disposable.fn()

beforeEach(() => {
restore()
})

afterEach(() => {
disposable.dispose()
})

describe('DvcViewer', () => {
it('should only be able to run a command once', async () => {
const mockConfig = {
getCliPath: () => 'sleep',
getPythonBinPath: () => undefined
} as Config
const dvcViewer = disposable.track(new DvcViewer(mockConfig))

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const getProcess = (dvcViewer: any, id: string): ViewableCliProcess =>
dvcViewer.processes[id]

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const createProcessSpy = spy(dvcViewer as any, 'createProcess')

await dvcViewer.run('command', dvcDemoPath, '10000')

expect(createProcessSpy).to.be.called
const viewableProcess = getProcess(
dvcViewer,
[dvcDemoPath, '10000'].join(':')
)
const showProcessSpy = spy(viewableProcess, 'show')

createProcessSpy.resetHistory()

await dvcViewer.run('command', dvcDemoPath, '10000')

expect(createProcessSpy).not.to.be.called
expect(showProcessSpy).to.be.calledOnce
})
})
})
2 changes: 1 addition & 1 deletion extension/src/vscode/pseudoTerminal.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ export class PseudoTerminal extends Disposable {
constructor(
processOutput: EventEmitter<string>,
processTerminated: EventEmitter<void>,
termName = 'DVC'
termName: string
) {
super()

Expand Down

0 comments on commit 616c35e

Please sign in to comment.