1# dockerpty: tty.py
2#
3# Copyright 2014 Chris Corbyn <chris@w3style.co.uk>
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17from __future__ import absolute_import
18
19import os
20import termios
21import tty
22import fcntl
23import struct
24
25
26def size(fd):
27    """
28    Return a tuple (rows,cols) representing the size of the TTY `fd`.
29
30    The provided file descriptor should be the stdout stream of the TTY.
31
32    If the TTY size cannot be determined, returns None.
33    """
34
35    if not os.isatty(fd.fileno()):
36        return None
37
38    try:
39        dims = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, 'hhhh'))
40    except:
41        try:
42            dims = (os.environ['LINES'], os.environ['COLUMNS'])
43        except:
44            return None
45
46    return dims
47
48
49class Terminal(object):
50    """
51    Terminal provides wrapper functionality to temporarily make the tty raw.
52
53    This is useful when streaming data from a pseudo-terminal into the tty.
54
55    Example:
56
57        with Terminal(sys.stdin, raw=True):
58            do_things_in_raw_mode()
59    """
60
61    def __init__(self, fd, raw=True):
62        """
63        Initialize a terminal for the tty with stdin attached to `fd`.
64
65        Initializing the Terminal has no immediate side effects. The `start()`
66        method must be invoked, or `with raw_terminal:` used before the
67        terminal is affected.
68        """
69
70        self.fd = fd
71        self.raw = raw
72        self.original_attributes = None
73
74
75    def __enter__(self):
76        """
77        Invoked when a `with` block is first entered.
78        """
79
80        self.start()
81        return self
82
83
84    def __exit__(self, *_):
85        """
86        Invoked when a `with` block is finished.
87        """
88
89        self.stop()
90
91
92    def israw(self):
93        """
94        Returns True if the TTY should operate in raw mode.
95        """
96
97        return self.raw
98
99
100    def start(self):
101        """
102        Saves the current terminal attributes and makes the tty raw.
103
104        This method returns None immediately.
105        """
106
107        if os.isatty(self.fd.fileno()) and self.israw():
108            self.original_attributes = termios.tcgetattr(self.fd)
109            tty.setraw(self.fd)
110
111
112    def stop(self):
113        """
114        Restores the terminal attributes back to before setting raw mode.
115
116        If the raw terminal was not started, does nothing.
117        """
118
119        if self.original_attributes is not None:
120            termios.tcsetattr(
121                self.fd,
122                termios.TCSADRAIN,
123                self.original_attributes,
124            )
125
126    def __repr__(self):
127        return "{cls}({fd}, raw={raw})".format(
128            cls=type(self).__name__,
129            fd=self.fd,
130            raw=self.raw)
131