diff --git a/src/segments/dvc.go b/src/segments/dvc.go new file mode 100644 index 000000000000..20f8b480d50b --- /dev/null +++ b/src/segments/dvc.go @@ -0,0 +1,81 @@ +package segments + +import "strings" + +// DvcStatus represents part of the status of a DVC repository +type DvcStatus struct { + ScmStatus +} + +func (s *DvcStatus) add(code string) { + switch code { + case "not in cache": + s.Missing++ + case "deleted": + s.Deleted++ + case "new": + s.Added++ + case "modified": + s.Modified++ + } +} + +const ( + DVCCOMMAND = "dvc" +) + +type Dvc struct { + Status *DvcStatus + Scm +} + +func (d *Dvc) Template() string { + return " \ue8d1 {{.Status.String}} " +} + +func (d *Dvc) Enabled() bool { + if !d.hasCommand(DVCCOMMAND) { + return false + } + + // Check if we're in a DVC repository + _, err := d.env.HasParentFilePath(".dvc", false) + if err != nil { + return false + } + + // run dvc status command + output, err := d.env.RunCommand(d.command, "status", "-q") + if err != nil { + return false + } + + statusFormats := d.options.KeyValueMap(StatusFormats, map[string]string{}) + d.Status = &DvcStatus{ScmStatus: ScmStatus{Formats: statusFormats}} + + if output == "" { + return true + } + + lines := strings.SplitSeq(output, "\n") + + for line := range lines { + if line == "" { + continue + } + + // DVC status output format: + // data.xml: modified + // or + // data/: not in cache + parts := strings.SplitN(line, ":", 2) + if len(parts) < 2 { + continue + } + + status := strings.TrimSpace(parts[1]) + d.Status.add(status) + } + + return true +} diff --git a/src/segments/dvc_test.go b/src/segments/dvc_test.go new file mode 100644 index 000000000000..f759397aeb55 --- /dev/null +++ b/src/segments/dvc_test.go @@ -0,0 +1,118 @@ +package segments + +import ( + "errors" + "fmt" + "testing" + + "github.com/jandedobbeleer/oh-my-posh/src/runtime" + "github.com/jandedobbeleer/oh-my-posh/src/runtime/mock" + "github.com/jandedobbeleer/oh-my-posh/src/segments/options" + + "github.com/stretchr/testify/assert" +) + +func TestDvcStatus(t *testing.T) { + cases := []struct { + Case string + Output string + OutputError error + HasCommand bool + HasDvcDir bool + ExpectedStatus string + ExpectedDisabled bool + }{ + { + Case: "not installed", + HasCommand: false, + ExpectedDisabled: true, + }, + { + Case: "not in DVC repo", + HasCommand: true, + HasDvcDir: false, + ExpectedDisabled: true, + }, + { + Case: "command error", + HasCommand: true, + HasDvcDir: true, + OutputError: fmt.Errorf("error"), + ExpectedDisabled: true, + }, + { + Case: "clean status", + HasCommand: true, + HasDvcDir: true, + Output: "", + ExpectedStatus: "", + }, + { + Case: "modified files", + HasCommand: true, + HasDvcDir: true, + Output: `data.xml: modified +model.pkl: modified`, + ExpectedStatus: "~2", + }, + { + Case: "new files", + HasCommand: true, + HasDvcDir: true, + Output: `data/new.csv: new +data/test.csv: new`, + ExpectedStatus: "+2", + }, + { + Case: "deleted files", + HasCommand: true, + HasDvcDir: true, + Output: `data.xml: deleted`, + ExpectedStatus: "-1", + }, + { + Case: "not in cache", + HasCommand: true, + HasDvcDir: true, + Output: `data/: not in cache`, + ExpectedStatus: "!1", + }, + { + Case: "mixed status", + HasCommand: true, + HasDvcDir: true, + Output: `data.xml: modified +model.pkl: new +old_data.csv: deleted +cache_data/: not in cache`, + ExpectedStatus: "+1 ~1 -1 !1", + }, + } + + for _, tc := range cases { + env := new(mock.Environment) + env.On("GOOS").Return("unix") + env.On("IsWsl").Return(false) + env.On("InWSLSharedDrive").Return(false) + env.On("HasCommand", DVCCOMMAND).Return(tc.HasCommand) + + if tc.HasDvcDir { + env.On("HasParentFilePath", ".dvc", false).Return(&runtime.FileInfo{}, nil) + } else { + env.On("HasParentFilePath", ".dvc", false).Return(&runtime.FileInfo{}, errors.New("not found")) + } + + env.On("RunCommand", DVCCOMMAND, []string{"status", "-q"}).Return(tc.Output, tc.OutputError) + + d := &Dvc{} + d.Init(options.Map{}, env) + + got := d.Enabled() + + assert.Equal(t, !tc.ExpectedDisabled, got, tc.Case) + if tc.ExpectedDisabled { + continue + } + assert.Equal(t, tc.ExpectedStatus, d.Status.String(), tc.Case) + } +} diff --git a/website/docs/segments/scm/dvc.mdx b/website/docs/segments/scm/dvc.mdx new file mode 100644 index 000000000000..0975b3481cbe --- /dev/null +++ b/website/docs/segments/scm/dvc.mdx @@ -0,0 +1,67 @@ +--- +id: dvc +title: DVC +sidebar_label: DVC +--- + +## What + +Display [DVC][dvc] (Data Version Control) information when in a DVC repository. + +## Sample Configuration + +import Config from '@site/src/components/Config.js'; + + + +## Options + +| Name | Type | Default | Description | +| ----------------- | :-----------------: | :-----: | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `native_fallback` | `boolean` | `false` | when set to `true` and `dvc.exe` is not available when inside a WSL2 shared Windows drive, we will fallback to the native `dvc` executable to fetch data. Not all information can be displayed in this case | +| `status_formats` | `map[string]string` | | a key, value map allowing to override how individual status items are displayed. For example, `"status_formats": { "Added": "Added: %d" }` will display the added count as `Added: 1` instead of `+1`. See the [Status](#status) section for available overrides | + +## Template ([info][templates]) + +:::note default template + +```template + \ue8d1 {{.Status.String}} +``` + +::: + +### Properties + +| Name | Type | Description | +| --------- | ----------- | ----------------------------------- | +| `.Status` | `DvcStatus` | changes in the worktree (see below) | + +### Status + +| Name | Type | Description | +| ------------ | --------- | ------------------------------------------ | +| `.Added` | `int` | number of new files | +| `.Modified` | `int` | number of modified files | +| `.Deleted` | `int` | number of deleted files | +| `.Missing` | `int` | number of files not in cache | +| `.Changed` | `boolean` | if the status contains changes or not | +| `.String` | `string` | a string representation of the changes above | + +Local changes use the following syntax: + +| Icon | Description | +| ---- | ----------- | +| `+` | added | +| `~` | modified | +| `-` | deleted | +| `!` | not in cache (missing) | + +[dvc]: https://dvc.org +[templates]: /configuration/templates.mdx