Moved all repo stats code from git.py to the base

powerline_shell_base.py

@@ -90,6 +90,65 @@ def draw_segment(self, idx):
             self.fgcolor(segment[4]),
             segment[3]))
 
+
+class RepoStats:
+    symbols = {
+        'detached': u'\u2693',
+        'ahead': u'\u2B06',
+        'behind': u'\u2B07',
+        'staged': u'\u2714',
+        'not_staged': u'\u270E',
+        'untracked': u'\u2753',
+        'conflicted': u'\u273C'
+    }
+
+    def __init__(self):
+        self.ahead = 0
+        self.behind = 0
+        self.untracked = 0
+        self.not_staged = 0
+        self.staged = 0
+        self.conflicted = 0
+
+    @property
+    def dirty(self):
+        qualifiers = [
+            self.untracked,
+            self.not_staged,
+            self.staged,
+            self.conflicted,
+        ]
+        return sum(qualifiers) > 0
+
+    def __getitem__(self, _key):
+        return getattr(self, _key)
+
+    def n_or_empty(self, _key):
+        """Given a string name of one of the properties of this class, returns
+        the value of the property as a string when the value is greater than
+        1. When it is not greater than one, returns an empty string.
+
+        As an example, if you want to show an icon for untracked files, but you
+        only want a number to appear next to the icon when there are more than
+        one untracked files, you can do:
+
+            segment = repo_stats.n_or_empty("untracked") + icon_string
+        """
+        return unicode(self[_key]) if int(self[_key]) > 1 else u''
+
+    def add_to_powerline(self, powerline, color):
+        def add(_key, fg, bg):
+            if self[_key]:
+                s = u" {}{} ".format(self.n_or_empty(_key), self.symbols[_key])
+                powerline.append(s, fg, bg)
+        add('ahead', color.GIT_AHEAD_FG, color.GIT_AHEAD_BG)
+        add('behind', color.GIT_BEHIND_FG, color.GIT_BEHIND_BG)
+        add('staged', color.GIT_STAGED_FG, color.GIT_STAGED_BG)
+        add('not_staged', color.GIT_NOTSTAGED_FG, color.GIT_NOTSTAGED_BG)
+        add('untracked', color.GIT_UNTRACKED_FG, color.GIT_UNTRACKED_BG)
+        add('conflicted', color.GIT_CONFLICTED_FG, color.GIT_CONFLICTED_BG)
+
+
 def get_valid_cwd():
     """ We check if the current working directory is valid or not. Typically
         happens when you checkout a different branch on git that doesn't have

segments/git.py

@@ -2,16 +2,6 @@
 import subprocess
 import os
 
-GIT_SYMBOLS = {
-    'detached': u'\u2693',
-    'ahead': u'\u2B06',
-    'behind': u'\u2B07',
-    'staged': u'\u2714',
-    'notstaged': u'\u270E',
-    'untracked': u'\u2753',
-    'conflicted': u'\u273C'
-}
-
 def get_PATH():
     """Normally gets the PATH from the OS. This function exists to enable
     easily mocking the PATH in tests.
@@ -43,33 +33,29 @@ def _get_git_detached_branch():
                          env=git_subprocess_env())
     detached_ref = p.communicate()[0].decode("utf-8").rstrip('\n')
     if p.returncode == 0:
-        branch = u'{} {}'.format(GIT_SYMBOLS['detached'], detached_ref)
+        branch = u'{} {}'.format(RepoStats.symbols['detached'], detached_ref)
     else:
         branch = 'Big Bang'
     return branch
 
 
 def parse_git_stats(status):
-    stats = {'untracked': 0, 'notstaged': 0, 'staged': 0, 'conflicted': 0}
+    stats = RepoStats()
     for statusline in status[1:]:
         code = statusline[:2]
         if code == '??':
-            stats['untracked'] += 1
+            stats.untracked += 1
         elif code in ('DD', 'AU', 'UD', 'UA', 'DU', 'AA', 'UU'):
-            stats['conflicted'] += 1
+            stats.conflicted += 1
         else:
             if code[1] != ' ':
-                stats['notstaged'] += 1
+                stats.not_staged += 1
             if code[0] != ' ':
-                stats['staged'] += 1
+                stats.staged += 1
 
     return stats
 
 
-def _n_or_empty(_dict, _key):
-    return _dict[_key] if int(_dict[_key]) > 1 else u''
-
-
 def add_git_segment(powerline):
     try:
         p = subprocess.Popen(['git', 'status', '--porcelain', '-b'],
@@ -84,33 +70,21 @@ def add_git_segment(powerline):
         return
 
     status = pdata[0].decode("utf-8").splitlines()
-
-    branch_info = parse_git_branch_info(status)
     stats = parse_git_stats(status)
-    dirty = (True if sum(stats.values()) > 0 else False)
+    branch_info = parse_git_branch_info(status)
 
     if branch_info:
+        stats.ahead = branch_info["ahead"]
+        stats.behind = branch_info["behind"]
         branch = branch_info['local']
     else:
         branch = _get_git_detached_branch()
 
     bg = Color.REPO_CLEAN_BG
     fg = Color.REPO_CLEAN_FG
-    if dirty:
+    if stats.dirty:
         bg = Color.REPO_DIRTY_BG
         fg = Color.REPO_DIRTY_FG
 
     powerline.append(' %s ' % branch, fg, bg)
-
-    def _add(_dict, _key, fg, bg):
-        if _dict[_key]:
-            _str = u' {}{} '.format(_n_or_empty(_dict, _key), GIT_SYMBOLS[_key])
-            powerline.append(_str, fg, bg)
-
-    if branch_info:
-        _add(branch_info, 'ahead', Color.GIT_AHEAD_FG, Color.GIT_AHEAD_BG)
-        _add(branch_info, 'behind', Color.GIT_BEHIND_FG, Color.GIT_BEHIND_BG)
-    _add(stats, 'staged', Color.GIT_STAGED_FG, Color.GIT_STAGED_BG)
-    _add(stats, 'notstaged', Color.GIT_NOTSTAGED_FG, Color.GIT_NOTSTAGED_BG)
-    _add(stats, 'untracked', Color.GIT_UNTRACKED_FG, Color.GIT_UNTRACKED_BG)
-    _add(stats, 'conflicted', Color.GIT_CONFLICTED_FG, Color.GIT_CONFLICTED_BG)
+    stats.add_to_powerline(powerline, Color)

test/repo_stats_test.py

@@ -0,0 +1,25 @@
+import unittest
+import powerline_shell_base as p
+
+
+class RepoStatsTest(unittest.TestCase):
+
+    def setUp(self):
+        self.repo_stats = p.RepoStats()
+        self.repo_stats.not_staged = 1
+        self.repo_stats.conflicted = 4
+
+    def test_dirty(self):
+        self.assertTrue(self.repo_stats.dirty)
+
+    def test_simple(self):
+        self.assertEqual(self.repo_stats.untracked, 0)
+
+    def test_n_or_empty__empty(self):
+        self.assertEqual(self.repo_stats.n_or_empty("not_staged"), u"")
+
+    def test_n_or_empty__n(self):
+        self.assertEqual(self.repo_stats.n_or_empty("conflicted"), u"4")
+
+    def test_index(self):
+        self.assertEqual(self.repo_stats["not_staged"], 1)

test/segments_test/git_test.py

@@ -3,14 +3,17 @@
 import tempfile
 import shutil
 import sh
+import powerline_shell_base as p
 import segments.git as git
 
+git.Color = mock.MagicMock()
+git.RepoStats = p.RepoStats
+
 
 class GitTest(unittest.TestCase):
 
     def setUp(self):
         self.powerline = mock.MagicMock()
-        git.Color = mock.MagicMock()
 
         self.dirname = tempfile.mkdtemp()
         sh.cd(self.dirname)