Add IPython support to pyspark-shell.

Suggested by / based on code from @MLnick

3 files changed, 21 insertions, 8 deletions

diff --git a/pyspark/README b/pyspark/README
index 55490e1a83..461176de7d 100644
--- a/pyspark/README
+++ b/pyspark/README
@@ -38,6 +38,9 @@ interacting with Java processes.  It can be installed from
 https://github.com/bartdag/py4j; make sure to install a version that
 contains at least the commits through b7924aabe9.
 
+PySpark requires the `argparse` module, which is included in Python 2.7
+and is is available for Python 2.6 through `pip` or `easy_install`.
+
 PySpark uses the `PYTHONPATH` environment variable to search for Python
 classes; Py4J should be on this path, along with any libraries used by
 PySpark programs.  `PYTHONPATH` will be automatically shipped to worker
diff --git a/pyspark/pyspark/shell.py b/pyspark/pyspark/shell.py
index 7ef30894cb..7012884abc 100644
--- a/pyspark/pyspark/shell.py
+++ b/pyspark/pyspark/shell.py
@@ -1,21 +1,30 @@
 """
 An interactive shell.
 """
+import argparse  # argparse is avaiable for Python < 2.7 through easy_install.
 import code
 import sys
 
 from pyspark.context import SparkContext
 
 
-def main(master='local'):
+def main(master='local', ipython=False):
     sc = SparkContext(master, 'PySparkShell')
-    print "Spark context available as sc."
-    code.interact(local={'sc': sc})
+    user_ns = {'sc' : sc}
+    banner = "Spark context avaiable as sc."
+    if ipython:
+        import IPython
+        IPython.embed(user_ns=user_ns, banner2=banner)
+    else:
+        print banner
+        code.interact(local=user_ns)
 
 
 if __name__ == '__main__':
-    if len(sys.argv) > 1:
-        master = sys.argv[1]
-    else:
-        master = 'local'
-    main(master)
+    parser = argparse.ArgumentParser()
+    parser.add_argument("master", help="Spark master host (default='local')",
+                        nargs='?', type=str, default="local")
+    parser.add_argument("-i", "--ipython", help="Run IPython shell",
+                        action="store_true")
+    args = parser.parse_args()
+    main(args.master, args.ipython)
diff --git a/pyspark/requirements.txt b/pyspark/requirements.txt
index 48fa2ab105..2464ca0074 100644
--- a/pyspark/requirements.txt
+++ b/pyspark/requirements.txt
@@ -4,3 +4,4 @@
 # install Py4J from git once https://github.com/pypa/pip/pull/526 is merged.
 
 # git+git://github.com/bartdag/py4j.git@b7924aabe9c5e63f0a4d8bbd17019534c7ec014e
+argparse