Fix stuff.

1 files changed, 12 insertions, 21 deletions

diff --git a/ponysay.py b/ponysay.py
index a75d44c..0728666 100755
--- a/ponysay.py
+++ b/ponysay.py
@@ -28,17 +28,11 @@ INSTALLDIR = '/usr'
 
 
 '''
-The user's home directory
-'''
-HOME = os.environ['HOME']
-
-
-'''
 The directories where pony files are stored, ttyponies/ are used if the terminal is Linux VT (also known as TTY)
 '''
 ponydirs = []
-if os.environ['TERM'] == 'linux':  _ponydirs = [INSTALLDIR + '/share/ponysay/ttyponies/',  HOME + '/.local/share/ponysay/ttyponies/']
-else:                              _ponydirs = [INSTALLDIR + '/share/ponysay/ponies/',     HOME + '/.local/share/ponysay/ponies/'   ]
+if os.environ['TERM'] == 'linux':  _ponydirs = [INSTALLDIR + '/share/ponysay/ttyponies/',  os.environ['HOME'] + '/.local/share/ponysay/ttyponies/']
+else:                              _ponydirs = [INSTALLDIR + '/share/ponysay/ponies/',     os.environ['HOME'] + '/.local/share/ponysay/ponies/'   ]
 for ponydir in _ponydirs:
     if os.path.isdir(ponydir):
         ponydirs.append(ponydir)
@@ -48,7 +42,7 @@ for ponydir in _ponydirs:
 The directories where quotes files are stored
 '''
 quotedirs = []
-_quotedirs = [INSTALLDIR + '/share/ponysay/quotes/',  HOME + '/.local/share/ponysay/quotes/']
+_quotedirs = [INSTALLDIR + '/share/ponysay/quotes/',  os.environ['HOME'] + '/.local/share/ponysay/quotes/']
 for quotedir in _quotedirs:
     if os.path.isdir(quotedir):
         quotedirs.append(quotedir)
@@ -113,15 +107,19 @@ class ponysay():
     def __getponypath(self, names = None):
         ponies = {}
         
-        for name in names:
-            if os.path.isfile(name):
-                return name
+        if names != None:
+            for name in names:
+                if os.path.isfile(name):
+                    return name
         
         for ponydir in ponydirs:
             for ponyfile in os.listdir(ponydir):
                 ponies[ponyfile[:-5]] = ponydir + ponyfile
         
-        return ponies[names[random.randrange(0, len(names))]]
+        if names == None:
+            names = list(ponies.keys())
+        
+        return ponies[names[random.randrange(0, len(names) - 1)]]
     
     '''
     Returns a list with all (pony, quote file) pairs
@@ -276,15 +274,8 @@ class ponysay():
     
     
     def print_pony(self, args):
-        ponycount = 0
-        for ponydir in ponydirs:
-            ponycount = len(os.listdir(ponydir))
-        if ponycount == 0:
-            sys.stderr.write('All the ponies are missing! Call the Princess!')
-            exit(1);
-        
         if args.message == None:
-            msg = sys.stdin.read()
+            msg = sys.stdin.read().strip()
         else:
             msg = args.message