feat(memory): add use_available and memory_type

JanDeDobbeleer · Sep 4, 2021 · ffad394 · ffad394
1 parent 4305b3a
commit ffad394
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 8 deletions.
diff --git a/docs/docs/memory.md b/docs/docs/memory.md
@@ -6,7 +6,7 @@ sidebar_label: Memory
 
 ## Memory
 
-Display memory usage percentage.
+Display physical memory or swap usage percentage.
 
 ## Sample Configuration
 
@@ -27,3 +27,5 @@ Display memory usage percentage.
 ## Properties
 
 - precision: `int` - the number of decimal places to show - defaults to `0`
+- use_available: `boolean` - whether to use available or free memory on Linux - defaults to `true`
+- memory_type: `enum` - whether to show `physical` memory or `swap` memory - defaults to `physical`
diff --git a/src/go.mod b/src/go.mod
@@ -18,9 +18,9 @@ require (
 	github.com/gookit/config/v2 v2.0.25
 	github.com/gookit/goutil v0.3.15 // indirect
 	github.com/huandu/xstrings v1.3.2 // indirect
+	github.com/jan0660/memory v0.1.2
 	github.com/mitchellh/copystructure v1.2.0 // indirect
 	github.com/mitchellh/mapstructure v1.4.1
-	github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58
 	github.com/sergi/go-diff v1.2.0 // indirect
 	github.com/shirou/gopsutil v3.21.8+incompatible
 	github.com/smartystreets/goconvey v1.6.4 // indirect

diff --git a/src/go.sum b/src/go.sum
@@ -59,6 +59,8 @@ github.com/huandu/xstrings v1.3.2/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq
 github.com/imdario/mergo v0.3.11/go.mod h1:jmQim1M+e3UYxmgPu/WyfjB3N3VflVyUjjjwH0dnCYA=
 github.com/imdario/mergo v0.3.12 h1:b6R2BslTbIEToALKP7LxUvijTsNI9TAe80pLWN2g/HU=
 github.com/imdario/mergo v0.3.12/go.mod h1:jmQim1M+e3UYxmgPu/WyfjB3N3VflVyUjjjwH0dnCYA=
+github.com/jan0660/memory v0.1.2 h1:Dxc2p3DJIJ+w2D2/OSqg3rf+ZQZ/lcYxiSXfiCNGRnA=
+github.com/jan0660/memory v0.1.2/go.mod h1:I9fFv85Tlthw0cS4nXzY9kCtCpNLuYsVVIxhejK50Cg=
 github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
 github.com/json-iterator/go v1.1.11 h1:uVUAXhF2To8cbw/3xN3pxj6kk7TYKs98NIrTqPlMWAQ=
 github.com/json-iterator/go v1.1.11/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
@@ -90,8 +92,6 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.1 h1:9f412s+6RmYXLWZSEzVVgPGK7C2PphHj5RJrvfx9AWI=
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
-github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0=
-github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/sergi/go-diff v1.0.0/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=

diff --git a/src/segment_memory.go b/src/segment_memory.go
@@ -3,7 +3,7 @@ package main
 import (
 	"strconv"
 
-	mem "github.com/pbnjay/memory"
+	mem "github.com/jan0660/memory"
 )
 
 type memory struct {
@@ -14,10 +14,19 @@ type memory struct {
 }
 
 const (
+	// Precision number of decimal places to show
 	Precision Property = "precision"
+	// UseAvailable if available memory should be used instead of free on Linux
+	UseAvailable Property = "use_available"
+	// MemoryType either physical or swap
+	MemoryType Property = "memory_type"
 )
 
 func (n *memory) enabled() bool {
+	if n.TotalMemory == 0 || n.FreeMemory == 0 {
+		// failed to get memory information
+		return false
+	}
 	return true
 }
 
@@ -31,6 +40,15 @@ func (n *memory) string() string {
 func (n *memory) init(props *properties, env environmentInfo) {
 	n.props = props
 	n.env = env
-	n.TotalMemory = mem.TotalMemory()
-	n.FreeMemory = mem.FreeMemory()
+	if props.getString(MemoryType, "physical") == "physical" {
+		n.TotalMemory = mem.TotalMemory()
+		if props.getBool(UseAvailable, true) {
+			n.FreeMemory = mem.AvailableMemory()
+			return
+		}
+		n.FreeMemory = mem.FreeMemory()
+		return
+	}
+	n.TotalMemory = mem.TotalSwap()
+	n.FreeMemory = mem.FreeSwap()
 }
diff --git a/src/segment_memory_test.go b/src/segment_memory_test.go
@@ -10,11 +10,13 @@ func TestMemory(t *testing.T) {
 	cases := []struct {
 		Case           string
 		ExpectedString string
+		ExpectDisabled bool
 		Memory         memory
 		Precision      int
 	}{
 		{Case: "50", ExpectedString: "50", Memory: memory{FreeMemory: 50, TotalMemory: 100}},
 		{Case: "50.0", ExpectedString: "50.0", Memory: memory{FreeMemory: 50, TotalMemory: 100}, Precision: 1},
+		{Case: "not enabled", ExpectDisabled: true, Memory: memory{FreeMemory: 0, TotalMemory: 0}},
 	}
 
 	for _, tc := range cases {
@@ -24,6 +26,10 @@ func TestMemory(t *testing.T) {
 				Precision: tc.Precision,
 			},
 		}
-		assert.Equal(t, tc.ExpectedString, tc.Memory.string(), tc.Case)
+		if tc.ExpectDisabled {
+			assert.Equal(t, false, tc.Memory.enabled(), tc.Case)
+		} else {
+			assert.Equal(t, tc.ExpectedString, tc.Memory.string(), tc.Case)
+		}
 	}
 }
diff --git a/themes/schema.json b/themes/schema.json
@@ -1668,6 +1668,19 @@
                     "title": "Precision",
                     "description": "number of decimal places to show",
                     "default": 0
+                  },
+                  "use_available": {
+                    "type": "boolean",
+                    "title": "Use available memory",
+                    "description": "whether to use available or free memory on Linux",
+                    "default": true
+                  },
+                  "memory_type": {
+                    "type": "string",
+                    "enum": ["physical", "swap"],
+                    "title": "Memory type",
+                    "description": "whether to show physical or swap memory",
+                    "default": "physical"
                   }
                 }
               }