From 07352818f0d49f653fc7262e578b4211157e55ea Mon Sep 17 00:00:00 2001
From: andrzejnovak <novak5andrzej@gmail.com>
Date: Tue, 16 Feb 2021 11:29:34 +0100
Subject: [PATCH] feat: module to display current screen session

---
 README.md         |  6 +++---
 defaults.go       |  3 +++
 main.go           |  1 +
 segment-screen.go | 30 ++++++++++++++++++++++++++++++
 themes.go         |  3 +++
 5 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 segment-screen.go

diff --git a/README.md b/README.md
index 411ed59f..e210e91b 100644
--- a/README.md
+++ b/README.md
@@ -249,12 +249,12 @@ Usage of powerline-go:
          (default "patched")
   -modules string
          The list of modules to load, separated by ','
-         (valid choices: aws, bzr, cwd, docker, docker-context, dotenv, duration, exit, fossil, git, gitlite, goenv, hg, host, jobs, kube, load, newline, nix-shell, node, perlbrew, perms, plenv, root, shell-var, shenv, ssh, svn, termtitle, terraform-workspace, time, user, venv, vgo)
+         (valid choices: aws, bzr, cwd, docker, docker-context, dotenv, duration, exit, fossil, git, gitlite, goenv, hg, host, jobs, kube, load, newline, nix-shell, node, perlbrew, perms, plenv, root, screen, shell-var, shenv, ssh, svn, termtitle, terraform-workspace, time, user, venv, vgo)
          Unrecognized modules will be invoked as 'powerline-go-MODULE' executable plugins and should output a (possibly empty) list of JSON objects that unmarshal to powerline-go's Segment structs.
          (default "venv,user,host,ssh,cwd,perms,git,hg,jobs,exit,root")
   -modules-right string
          The list of modules to load anchored to the right, for shells that support it, separated by ','
-         (valid choices: aws, bzr, cwd, docker, docker-context, dotenv, duration, exit, fossil, git, gitlite, goenv, hg, host, jobs, kube, load, newline, nix-shell, node, perlbrew, perms, plenv, root, shell-var, shenv, ssh, svn, termtitle, terraform-workspace, time, user, venv, vgo)
+         (valid choices: aws, bzr, cwd, docker, docker-context, dotenv, duration, exit, fossil, git, gitlite, goenv, hg, host, jobs, kube, load, newline, nix-shell, node, perlbrew, perms, plenv, root, screen, shell-var, shenv, ssh, svn, termtitle, terraform-workspace, time, user, venv, vgo)
          Unrecognized modules will be invoked as 'powerline-go-MODULE' executable plugins and should output a (possibly empty) list of JSON objects that unmarshal to powerline-go's Segment structs.
   -newline
          Show the prompt on a new line
@@ -267,7 +267,7 @@ Usage of powerline-go:
          Use '~' for your home dir. You may need to escape this character to avoid shell substitution.
   -priority string
          Segments sorted by priority, if not enough space exists, the least priorized segments are removed first. Separate with ','
-         (valid choices: aws, bzr, cwd, docker, docker-context, dotenv, duration, exit, fossil, git, gitlite, hg, host, jobs, kube, load, newline, nix-shell, node, perlbrew, perms, plenv, root, shell-var, shenv, ssh, svn, termtitle, terraform-workspace, time, user, venv, vgo)
+         (valid choices: aws, bzr, cwd, docker, docker-context, dotenv, duration, exit, fossil, git, gitlite, hg, host, jobs, kube, load, newline, nix-shell, node, perlbrew, perms, plenv, root, screen, shell-var, shenv, ssh, svn, termtitle, terraform-workspace, time, user, venv, vgo)
          (default "root,cwd,user,host,ssh,perms,git-branch,git-status,hg,jobs,exit,cwd-path")
   -shell string
          Set this to your shell type
diff --git a/defaults.go b/defaults.go
index 7c95c165..f8fb6e8a 100644
--- a/defaults.go
+++ b/defaults.go
@@ -233,6 +233,9 @@ var defaults = Config{
 			VirtualEnvFg: 00,
 			VirtualEnvBg: 35, // a mid-tone green
 
+			ScreenFg: 15, // white
+			ScreenBg: 26, // blue
+
 			VirtualGoFg: 220, // approx. Secondary Yellow
 			VirtualGoBg: 38,  // approx. Gopher Blue
 
diff --git a/main.go b/main.go
index 1a1a8c24..e07283dc 100644
--- a/main.go
+++ b/main.go
@@ -91,6 +91,7 @@ var modules = map[string]func(*powerline) []pwl.Segment{
 	"perms":               segmentPerms,
 	"rbenv":               segmentRbenv,
 	"root":                segmentRoot,
+	"screen":              segmentScreenName,
 	"shell-var":           segmentShellVar,
 	"shenv":               segmentShEnv,
 	"ssh":                 segmentSSH,
diff --git a/segment-screen.go b/segment-screen.go
new file mode 100644
index 00000000..fcf6405d
--- /dev/null
+++ b/segment-screen.go
@@ -0,0 +1,30 @@
+package main
+
+import (
+	"os"
+	"strings"
+
+	pwl "github.com/justjanne/powerline-go/powerline"
+)
+
+func segmentScreenName(p *powerline) []pwl.Segment {
+	var env string
+	if env == "" {
+		env, _ = os.LookupEnv("STY")
+	}
+	if env == "" {
+		return []pwl.Segment{}
+	}
+	// envName := env
+	envName := strings.Split(env, ".")[1]
+	if p.cfg.VenvNameSizeLimit > 0 && len(envName) > p.cfg.VenvNameSizeLimit {
+		envName = p.symbols.VenvIndicator
+	}
+
+	return []pwl.Segment{{
+		Name:       "screen",
+		Content:    envName,
+		Foreground: p.theme.ScreenFg,
+		Background: p.theme.ScreenBg,
+	}}
+}
diff --git a/themes.go b/themes.go
index 8f69df21..f864d964 100644
--- a/themes.go
+++ b/themes.go
@@ -114,6 +114,9 @@ type Theme struct {
 	VirtualEnvFg uint8
 	VirtualEnvBg uint8
 
+	ScreenFg uint8
+	ScreenBg uint8
+
 	VirtualGoFg uint8
 	VirtualGoBg uint8