feat: show python segment when Jupyter Notebook files are present

JanDeDobbeleer · Oct 8, 2020 · 0b0fd4c · 0b0fd4c
1 parent 59ed9c2
commit 0b0fd4c
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 6 deletions.
diff --git a/docs/docs/segment-python.md b/docs/docs/segment-python.md
@@ -6,7 +6,7 @@ sidebar_label: Python
 
 ## What
 
-Display the currently active python version and virtualenv when a folder contains `.py` files.
+Display the currently active python version and virtualenv when a folder contains `.py` files or `.ipynb` files.
 Supports conda, virtualenv and pyenv.
 
 ## Sample Configuration

diff --git a/segment_python.go b/segment_python.go
@@ -30,7 +30,7 @@ func (p *python) init(props *properties, env environmentInfo) {
 }
 
 func (p *python) enabled() bool {
-	if !p.env.hasFiles("*.py") {
+	if !p.env.hasFiles("*.py") && !p.env.hasFiles("*.ipynb") {
 		return false
 	}
 	pythonVersions := []string{

diff --git a/segment_python_test.go b/segment_python_test.go
@@ -15,7 +15,8 @@ type pythonArgs struct {
 	pathSeparator    string
 	pythonVersion    string
 	python3Version   string
-	hasFiles         bool
+	hasPyFiles         bool
+	hasNotebookFiles   bool
 }
 
 func newPythonArgs() *pythonArgs {
@@ -27,13 +28,15 @@ func newPythonArgs() *pythonArgs {
 		pathSeparator:    "/",
 		pythonVersion:    "",
 		python3Version:   "",
-		hasFiles:         true,
+		hasPyFiles:       true,
+		hasNotebookFiles: true,
 	}
 }
 
 func bootStrapPythonTest(args *pythonArgs) *python {
 	env := new(MockedEnvironment)
-	env.On("hasFiles", "*.py").Return(args.hasFiles)
+	env.On("hasFiles", "*.py").Return(args.hasPyFiles)
+	env.On("hasFiles", "*.ipynb").Return(args.hasNotebookFiles)
 	env.On("runCommand", "python", []string{"--version"}).Return(args.pythonVersion)
 	env.On("runCommand", "python3", []string{"--version"}).Return(args.python3Version)
 	env.On("getenv", "VIRTUAL_ENV").Return(args.virtualEnvName)
@@ -49,11 +52,48 @@ func bootStrapPythonTest(args *pythonArgs) *python {
 
 func TestPythonWriterDisabledNoPythonFiles(t *testing.T) {
 	args := newPythonArgs()
-	args.hasFiles = false
+	args.hasPyFiles = false
+	args.hasNotebookFiles = false
+	args.python3Version = "3.4.5"
 	python := bootStrapPythonTest(args)
 	assert.False(t, python.enabled(), "there are no Python files in the current folder")
 }
 
+func TestPythonWriterDisabledHasPythonFiles(t *testing.T) {
+	args := newPythonArgs()
+	args.hasPyFiles = true
+	args.hasNotebookFiles = false
+	args.python3Version = "3.4.5"
+	python := bootStrapPythonTest(args)
+	assert.True(t, python.enabled(), "there should be a Python file in the current folder")
+}
+
+func TestPythonWriterDisabledHasJupyterNotebookFiles(t *testing.T) {
+	args := newPythonArgs()
+	args.hasPyFiles = false
+	args.hasNotebookFiles = true
+	args.python3Version = "3.4.5"
+	python := bootStrapPythonTest(args)
+	assert.True(t, python.enabled(), "there should be a Jupyter Notebook file in the current folder")
+}
+
+func TestPythonWriterDisabledHasPyAndJupyterNotebookFiles(t *testing.T) {
+	args := newPythonArgs()
+	args.hasPyFiles = true
+	args.hasNotebookFiles = true
+	args.python3Version = "3.4.5"
+	python := bootStrapPythonTest(args)
+	assert.True(t, python.enabled(), "there should be a Jupyter Notebook file in the current folder")
+}
+
+func TestPythonWriterDisabledHasPyAndJupyterNotebookFilesButNoVersion(t *testing.T) {
+	args := newPythonArgs()
+	args.hasPyFiles = true
+	args.hasNotebookFiles = true
+	python := bootStrapPythonTest(args)
+	assert.False(t, python.enabled(), "there should be a Jupyter Notebook file in the current folder")
+}
+
 func TestPythonWriterDisabledNoPythonInstalled(t *testing.T) {
 	args := newPythonArgs()
 	python := bootStrapPythonTest(args)