Improved routine for terminal dimensions. Works on windows too. (#127)

tests/terminal_test.py

@@ -27,15 +27,11 @@ class TerminalTest(unittest.TestCase):
 
   def setUp(self):
     super(TerminalTest, self).setUp()
-    self.environ_orig = terminal.os.environ
-    self.open_orig = terminal.os.open
-    self.terminal_orig = terminal.TerminalSize
+    self._get_terminal_size_orig = terminal.shutil.get_terminal_size
 
   def tearDown(self):
     super(TerminalTest, self).tearDown()
-    terminal.os.environ = self.environ_orig
-    terminal.os.open = self.open_orig
-    terminal.TerminalSize = self.terminal_orig
+    terminal.shutil.get_terminal_size = self._get_terminal_size_orig
 
   def testAnsiCmd(self):
     self.assertEqual('\033[0m', terminal._AnsiCmd(['reset']))
@@ -66,22 +62,8 @@ def testEncloseAnsi(self):
     ansi_enclosed = '\001\033[5;32;44m\002ansi\001\033[0m\002 length'
     self.assertEqual(ansi_enclosed, terminal.EncloseAnsiText(ansi_text))
 
-  def testTerminalSize(self):
-    # pylint: disable=unused-argument
-    def StubOpen(args, *kwargs):
-      raise IOError
-    terminal.open = StubOpen
-    terminal.os.environ = {}
-    # Raise exceptions on ioctl and environ and assign a default.
-    self.assertEqual((24, 80), terminal.TerminalSize())
-    terminal.os.environ = {'LINES': 'bogus', 'COLUMNS': 'bogus'}
-    self.assertEqual((24, 80), terminal.TerminalSize())
-    # Still raise exception on ioctl and use environ.
-    terminal.os.environ = {'LINES': '10', 'COLUMNS': '20'}
-    self.assertEqual((10, 20), terminal.TerminalSize())
-
   def testLineWrap(self):
-    terminal.TerminalSize = lambda: (5, 11)
+    terminal.shutil.get_terminal_size = lambda: (11, 5)
     text = ''
     self.assertEqual(text, terminal.LineWrap(text))
     text = 'one line'
@@ -145,19 +127,15 @@ def setUp(self):
     sys.stdout = FakeTerminal()
     self.get_ch_orig = terminal.Pager._GetCh
     terminal.Pager._GetCh = lambda self: 'q'
-    self.ts_orig = terminal.TerminalSize
-    terminal.TerminalSize = lambda: (24, 80)
 
     self.p = terminal.Pager()
 
   def tearDown(self):
     super(PagerTest, self).tearDown()
     terminal.Pager._GetCh = self.get_ch_orig
-    terminal.TerminalSize = self.ts_orig
     sys.stdout = sys.__stdout__
 
   def testPager(self):
-    self.assertEqual(terminal.TerminalSize()[0], self.p._cli_lines)
 
     self.p.Clear()
     self.assertEqual('', self.p._text)
@@ -170,7 +148,7 @@ def testPage(self):
       txt += '%d a random line of text here\n' % i
     self.p._text = txt
     self.p.Page()
-    self.assertEqual(terminal.TerminalSize()[0]+2, sys.stdout.CountLines())
+    self.assertEqual(self.p._cli_lines+2, sys.stdout.CountLines())
 
     sys.stdout.output = ''
     self.p = terminal.Pager()

textfsm/terminal.py

@@ -18,9 +18,8 @@
 """Simple terminal related routines."""
 
 import getopt
-import os
 import re
-import struct
+import shutil
 import sys
 import time
 
@@ -32,8 +31,6 @@
 except (ImportError, ModuleNotFoundError):
   pass
 
-__version__ = '0.1.1'
-
 # ANSI, ISO/IEC 6429 escape sequences, SGR (Select Graphic Rendition) subset.
 SGR = {
     'reset': 0,
@@ -168,21 +165,6 @@ def EncloseAnsiText(text):
   return sgr_re.sub(lambda x: ANSI_START + x.group(1) + ANSI_END, text)
 
 
-def TerminalSize():
-  """Returns terminal length and width as a tuple."""
-  try:
-    with open(os.ctermid()) as tty_instance:
-      length_width = struct.unpack(
-          'hh', fcntl.ioctl(tty_instance.fileno(), termios.TIOCGWINSZ, '1234')
-      )
-  except (IOError, OSError, NameError):
-    try:
-      length_width = (int(os.environ['LINES']), int(os.environ['COLUMNS']))
-    except (ValueError, KeyError):
-      length_width = (24, 80)
-  return length_width
-
-
 def LineWrap(text, omit_sgr=False):
   """Break line to fit screen width, factoring in ANSI/SGR escape sequences.
 
@@ -194,7 +176,7 @@ def LineWrap(text, omit_sgr=False):
     Text with additional line wraps inserted for lines grater than the width.
   """
 
-  def _SplitWithSgr(text_line):
+  def _SplitWithSgr(text_line, width):
     """Tokenise the line so that the sgr sequences can be omitted."""
     token_list = sgr_re.split(text_line)
     text_line_list = []
@@ -226,20 +208,20 @@ def _SplitWithSgr(text_line):
 
   # We don't use textwrap library here as it insists on removing
   # trailing/leading whitespace (pre 2.6).
-  (_, width) = TerminalSize()
+  (term_width, _) = shutil.get_terminal_size()
   text = str(text)
   text_multiline = []
   for text_line in text.splitlines():
     # Is this a line that needs splitting?
-    while (omit_sgr and (len(StripAnsiText(text_line)) > width)) or (
-        len(text_line) > width
+    while (omit_sgr and (len(StripAnsiText(text_line)) > term_width)) or (
+        len(text_line) > term_width
     ):
       # If there are no sgr escape characters then do a straight split.
       if not omit_sgr:
-        text_multiline.append(text_line[:width])
-        text_line = text_line[width:]
+        text_multiline.append(text_line[:term_width])
+        text_line = text_line[term_width:]
       else:
-        (multiline_line, text_line) = _SplitWithSgr(text_line)
+        (multiline_line, text_line) = _SplitWithSgr(text_line, term_width)
         text_multiline.append(multiline_line)
     if text_line:
       text_multiline.append(text_line)
@@ -318,7 +300,7 @@ def SetLines(self, lines):
       ValueError, TypeError: Not a valid integer representation.
     """
 
-    (self._cli_lines, self._cli_cols) = TerminalSize()
+    (self._cli_cols, self._cli_lines) = shutil.get_terminal_size()
 
     if lines:
       self._cli_lines = int(lines)
@@ -472,7 +454,8 @@ def main(argv=None):
     # Prints the size of the terminal and returns.
     # Mutually exclusive to the paging of text and overrides that behaviour.
     if opt in ('-s', '--size'):
-      print('Length: %d, Width: %d' % TerminalSize())
+      print(
+        'Width: %d, Length: %d' % shutil.get_terminal_size())
       return 0
     elif opt in ('-d', '--delay'):
       isdelay = True