Python script that polls memory usage for a Linux user

watch_mem_usage.py

#!/usr/bin/python
"""
This script is intended to run infinitely (ctrl-c to exit). It executes a "ps" command, and then 
parses the result to determine the total memory usage for a given user account. It will output a 
warning message if that value exceeds a specified threshold.

The following values can be configured:

USER - the account to be watched. The $USER bash environment variable is the default.
MEM_WARNING_LEVEL - warning message will be shown if current usage is above this
DELAY - the polling interval for executing the "ps" command
"""

import subprocess
from datetime import datetime
import time
import sys

# see doc @ top for information on these values
MEM_WARNING_LEVEL = 500000
USER = '$USER'
DELAY = 3
PS_COMMAND = "ps -u %(user)s -o rss,etime,pid,command | awk '{ SUM += $1 } END { print SUM }';" % dict(user=USER)


def hilite(string, status, bold):
  """ Colors a string of text for shell output.
  
  param: string   the text to be marked-up
  param: status   True=green / False=red / integer-value for a specific color code
  param: bold     True/False
  
  return: marked-up string of text
  """
  if not sys.stdout.isatty():
    return string

  attr = []
  if status is True:
    # green
    attr.append('32')
  elif status is False:
    # red
    attr.append('31')
  elif int(status) > 0:
    attr.append(status)

  if bold:
    attr.append('1')
  return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), string)


def make_output(text):
  """ Assembles the content to be output for each cycle. This consists of a timestamp, followed
  by the amount of memory being used. If that value is greater than the MEM_WARNING_LEVEL value,
  the message is prepended with a warning/alert, in red text.
  
  param: text   the text to be evaluated, which should simply be the memory usage value, from the "ps" cmd.
  
  return: text to be outputted to the shell
  """
  msgs = list()

  msg = text[:-1] # trim off trailing \r\n
  time_str = datetime.now().strftime("%Y-%d-%m, %H:%M:%S")

  if int(msg) > MEM_WARNING_LEVEL:
    msgs.append(hilite("%s: *** Warning: mem usage is above %s ***" % (time_str, MEM_WARNING_LEVEL), False, True))

  hl_msg = hilite(msg, '33', False)
  msgs.append("%s: Current mem usage: %s" % (hilite(time_str, True, True), hl_msg))
  return "\r\n".join(msgs)


def main():
  """ Executes the PS_COMMAND, then sleeps for DELAY seconds, and then repeats. """
  for (output, err) in iter(lambda: subprocess.Popen([PS_COMMAND], stdout=subprocess.PIPE, shell=True).communicate(), False):
    print(make_output(output))
    time.sleep(DELAY)


if __name__ == '__main__':
  try:
    main()
  except KeyboardInterrupt: # ctrl-c exits cleanly
    sys.exit(0)