src/org/apache/pig/scripting/jython/JythonScriptEngine.java - pig - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 package org.apache.pig.scripting.jython;

 import java.io.BufferedReader;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.regex.Pattern;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.pig.FuncSpec;
 import org.apache.pig.PigServer;
 import org.apache.pig.backend.executionengine.ExecException;
 import org.apache.pig.impl.PigContext;
 import org.apache.pig.impl.util.ObjectSerializer;
 import org.apache.pig.impl.util.Utils;
 import org.apache.pig.parser.ParserException;
 import org.apache.pig.scripting.ScriptEngine;
 import org.apache.pig.tools.pigstats.PigStats;
 import org.python.core.ClasspathPyImporter;
 import org.python.core.Py;
 import org.python.core.PyException;
 import org.python.core.PyFrame;
 import org.python.core.PyFunction;
 import org.python.core.PyInteger;
 import org.python.core.PyJavaPackage;
 import org.python.core.PyObject;
 import org.python.core.PyString;
 import org.python.core.PyStringMap;
 import org.python.core.PySystemState;
 import org.python.core.PyTuple;
 import org.python.modules.zipimport.zipimporter;
 import org.python.util.PythonInterpreter;

 /**
  * Implementation of the script engine for Jython
  */
 public class JythonScriptEngine extends ScriptEngine {
     private static final Log LOG = LogFactory.getLog(JythonScriptEngine.class);

     /**
      * Language Interpreter Uses static holder pattern
      */
     private static class Interpreter {
         static final PythonInterpreter interpreter;
         static final ArrayList<String> filesLoaded = new ArrayList<String>();
         static final String JVM_JAR;

         static {
             // should look like: file:JVM_JAR!/java/lang/Object.class
             String rpath = Object.class.getResource("Object.class").getPath();
             JVM_JAR = rpath.replaceAll("^file:(.*)!/java/lang/Object.class$", "$1");

             // Determine if a usable python.cachedir has been provided
             // if not, certain uses of jython's import will not work e.g., so create a tmp dir
             //  - from some.package import *
             //  - import non.jvm.package
             try {
                 String skip = System.getProperty(PySystemState.PYTHON_CACHEDIR_SKIP, "false");
                 if (skip.equalsIgnoreCase("true")) {
                     LOG.warn("jython cachedir skipped, jython may not work");
                 } else {
                     File tmp = null;
                     String cdir = System.getProperty(PySystemState.PYTHON_CACHEDIR);
                     if (cdir != null) {
                         tmp = new File(cdir);
                         if (tmp.canWrite() == false) {
                             LOG.error("CACHEDIR: not writable");
                             throw new RuntimeException("python.cachedir not writable: " + cdir);
                         }
                     }
                     if (tmp == null) {
                         tmp = File.createTempFile("pig_jython_", "");
                         tmp.delete();
                         if (tmp.mkdirs() == false) {
                             LOG.warn("unable to create a tmp dir for the cache, jython may not work");
                         } else {
                             LOG.info("created tmp python.cachedir=" + tmp);
                             System.setProperty(PySystemState.PYTHON_CACHEDIR, tmp.getAbsolutePath());
                         }
                         Runtime.getRuntime().addShutdownHook(new DirDeleter(tmp));
                     }
                 }
                 // local file system import path elements: current dir, JYTHON_HOME/Lib
                 Py.getSystemState().path.append(new PyString(System.getProperty("user.dir")));
                 String jyhome = System.getenv("JYTHON_HOME");
                 if (jyhome != null) {
                     Py.getSystemState().path.append(new PyString(jyhome + File.separator + "Lib"));
                 }
             } catch (Exception e) {
                 LOG.warn("issue with jython cache dir", e);
             }

             // cacdedir now configured, allocate the python interpreter
             interpreter = new PythonInterpreter();
         }

         /**
          * ensure the decorator functions are defined in the interpreter, and
          * manage the module import dependencies.
          * @param path       location of a file to exec in the interpreter
          * @param pigContext if non-null, module import state is tracked
          * @throws IOException
          */
         static synchronized void init(String path, PigContext pigContext) throws IOException {
             // Decorators -
             // "schemaFunction"
             // "outputSchema"
             // "outputSchemaFunction"

             if (!filesLoaded.contains(path)) {
                 // attempt addition of schema decorator handler, fail silently
                 interpreter.exec("def outputSchema(schema_def):\n"
                         + "    def decorator(func):\n"
                         + "        func.outputSchema = schema_def\n"
                         + "        return func\n"
                         + "    return decorator\n\n");

                 interpreter.exec("def outputSchemaFunction(schema_def):\n"
                         + "    def decorator(func):\n"
                         + "        func.outputSchemaFunction = schema_def\n"
                         + "        return func\n"
                         + "    return decorator\n");

                 interpreter.exec("def schemaFunction(schema_def):\n"
                         + "     def decorator(func):\n"
                         + "         func.schemaFunction = schema_def\n"
                         + "         return func\n"
                         + "     return decorator\n\n");

                 InputStream is = getScriptAsStream(path);
                 if (is == null) {
                     throw new IllegalStateException("unable to create a stream for path: " + path);
                 }
                 try {
                     execfile(is, path, pigContext);
                 } finally {
                     is.close();
                 }
             }
         }

         /**
          * does not call script.close()
          * @param script
          * @param path
          * @param pigContext
          * @throws ExecException
          */
         static void execfile(InputStream script, String path, PigContext pigContext) throws ExecException {
             try {

                 if( pigContext != null ) {
                   String [] argv;
                   try {
                       argv = (String[])ObjectSerializer.deserialize(
                               pigContext.getProperties().getProperty(PigContext.PIG_CMD_ARGS_REMAINDERS));
                   } catch (IOException e) {
                       throw new ExecException("Cannot deserialize command line arguments", e);
                   }
                   PySystemState  state = Py.getSystemState();
                   state.argv.clear();
                   if( argv != null ) {
                     for (String str : argv ) {
                       state.argv.append(new PyString(str));
                     }
                   } else {
                     LOG.warn(PigContext.PIG_CMD_ARGS_REMAINDERS
                       + " is empty. This is not expected unless on testing." );
                   }
                 }

                 // determine the current module state
                 Map<String, String> before = pigContext != null ? getModuleState() : null;
                 if (before != null) {
                     // os.py, stax.py and posixpath.py are part of the initial state
                     // if Lib directory is present and without including os.py, modules
                     // which import os fail
                     Set<String> includePyModules = new HashSet<String>();
                     for (String key : before.keySet()) {
                         // $py.class is created if Lib folder is writable
                         if (key.endsWith(".py") || key.endsWith("$py.class")) {
                             includePyModules.add(key);
                         }
                     }
                     before.keySet().removeAll(includePyModules);
                 }

                 // exec the code, arbitrary imports are processed
                 interpreter.execfile(script, path);

                 // determine the 'post import' module state
                 Map<String, String> after = pigContext != null ? getModuleState() : null;

                 // add the module files to the context
                 if (after != null && pigContext != null) {
                     after.keySet().removeAll(before.keySet());
                     for (Map.Entry<String, String> entry : after.entrySet()) {
                         String modulename = entry.getKey();
                         String modulepath = entry.getValue();
                         if (modulepath.equals(JVM_JAR)) {
                             continue;
                         } else if (modulepath.endsWith(".jar") || modulepath.endsWith(".zip")) {
                             pigContext.addScriptJar(modulepath);
                         } else {
                             pigContext.addScriptFile(modulename, modulepath);
                         }
                     }
                 }
             } catch (PyException e) {
                 if (e.match(Py.SystemExit)) {
                     PyObject value = e.value;
                     if (PyException.isExceptionInstance(e.value)) {
                         value = value.__findattr__("code");
                     }
                     if (new  PyInteger(0).equals(value)) {
                         LOG.info("Script invoked sys.exit(0)");
                         return;
                     }
                 }
                 String message = "Python Error. " + e;
                 throw new ExecException(message, 1121, e);
             }
         }

         static String get(String name) {
             return interpreter.get(name).toString();
         }

         static void setMain(boolean isMain) {
             if (isMain) {
                 interpreter.set("__name__", "__main__");
             } else {
                 interpreter.set("__name__", "__lib__");
             }
         }

         /**
          * get the state of modules currently loaded
          * @return a map of module name to module file (absolute path)
          */
         private static Map<String, String> getModuleState() {
             // determine the current module state
             Map<String, String> files = new HashMap<String, String>();
             PyStringMap modules = (PyStringMap) Py.getSystemState().modules;
             for (PyObject kvp : modules.iteritems().asIterable()) {
                 PyTuple tuple = (PyTuple) kvp;
                 String name = tuple.get(0).toString();
                 Object value = tuple.get(1);
                 // inspect the module to determine file location and status
                 try {
                     Object fileEntry = null;
                     Object loader = null;
                     if (value instanceof PyJavaPackage ) {
                         fileEntry = ((PyJavaPackage) value).__file__;
                     } else if (value instanceof PyObject) {
                         // resolved through the filesystem (or built-in)
                         PyObject dict = ((PyObject) value).getDict();
                         if (dict != null) {
                             fileEntry = dict.__finditem__("__file__");
                             loader = dict.__finditem__("__loader__");
                         } // else built-in
                     }   // else some system module?

                     if (fileEntry != null) {
                         File file = resolvePyModulePath(fileEntry.toString(), loader);
                         if (file.exists()) {
                             String apath = file.getAbsolutePath();
                             if (apath.endsWith(".jar") || apath.endsWith(".zip")) {
                                 // jar files are simple added to the pigContext
                                 files.put(apath, apath);
                             } else {
                                 // determine the relative path that the file should have in the jar
                                 int pos = apath.lastIndexOf(File.separatorChar + name.replace('.', File.separatorChar));
                                 if (pos > 0) {
                                     files.put(apath.substring(pos + 1), apath);
                                 } else {
                                     files.put(apath, apath);
                                 }
                             }
                         } else {
                             LOG.warn("module file does not exist: " + name + ", " + file);
                         }
                     } // else built-in
                 } catch (Exception e) {
                     LOG.warn("exception while retrieving module state: " + value, e);
                 }
             }
             return files;
         }
     }

     private static File resolvePyModulePath(String path, Object loader) {
         File file = new File(path);
         if (!file.exists() && loader != null) {
             if(path.startsWith(ClasspathPyImporter.PYCLASSPATH_PREFIX) && loader instanceof ClasspathPyImporter) {
                 path = path.replaceFirst(ClasspathPyImporter.PYCLASSPATH_PREFIX, "");
                 URL resource = ScriptEngine.class.getResource(path);
                 if (resource == null) {
                     resource = ScriptEngine.class.getResource(File.separator + path);
                 }
                 if (resource != null) {
                     return new File(resource.getFile());
                 }
             } else if (loader instanceof zipimporter) {
                 zipimporter importer = (zipimporter) loader;
                 return new File(importer.archive);
             } //JavaImporter??
         }
         return file;
     }

     @Override
     public void registerFunctions(String path, String namespace, PigContext pigContext)
     throws IOException {
         Interpreter.setMain(false);
         Interpreter.init(path, pigContext);
         pigContext.addScriptJar(getJarPath(PythonInterpreter.class));
         PythonInterpreter pi = Interpreter.interpreter;
         @SuppressWarnings("unchecked")
         List<PyTuple> locals = ((PyStringMap) pi.getLocals()).items();
         namespace = (namespace == null) ? "" : namespace + NAMESPACE_SEPARATOR;
         try {
             for (PyTuple item : locals) {
                 String key = (String) item.get(0);
                 Object value = item.get(1);
                 FuncSpec funcspec = null;
                 if (!key.startsWith("__") && !key.equals("schemaFunction")
                         && !key.equals("outputSchema")
                         && !key.equals("outputSchemaFunction")
                         && (value instanceof PyFunction)
                         && (((PyFunction)value).__findattr__("schemaFunction")== null)) {
                     PyObject obj = ((PyFunction)value).__findattr__("outputSchema");
                     if(obj != null) {
                         Utils.getSchemaFromString(obj.toString());
                     }
                     funcspec = new FuncSpec(JythonFunction.class.getCanonicalName() + "('"
                             + path + "','" + key +"')");
                     pigContext.registerFunction(namespace + key, funcspec);
                     LOG.info("Register scripting UDF: " + namespace + key);
                 }
             }
         } catch (ParserException pe) {
             throw new IOException(
                     "Error parsing schema for script function from the decorator",
                     pe);
         }
         pigContext.addScriptFile(path);
         Interpreter.setMain(true);
     }

     /**
      * Gets the Python function object.
      * @param path Path of the jython script file containing the function.
      * @param functionName Name of the function
      * @return a function object
      * @throws IOException
      */
     public static PyFunction getFunction(String path, String functionName) throws IOException {
         Interpreter.setMain(false);
         Interpreter.init(path, null);
         return (PyFunction) Interpreter.interpreter.get(functionName);
     }

     @Override
     protected Map<String, List<PigStats>> main(PigContext pigContext, String scriptFile)
             throws IOException {
         if (System.getProperty(PySystemState.PYTHON_CACHEDIR_SKIP)==null)
             System.setProperty(PySystemState.PYTHON_CACHEDIR_SKIP, "false");

         PigServer pigServer = new PigServer(pigContext, false);

         // register dependencies
         String jythonJar = getJarPath(PythonInterpreter.class);
         if (jythonJar != null) {
             pigServer.registerJar(jythonJar);
         }

         File f = new File(scriptFile);

         if (!f.canRead()) {
             throw new IOException("Can't read file: " + scriptFile);
         }

         FileInputStream fis1 = new FileInputStream(scriptFile);
         try {
             if (hasFunction(fis1)) {
                 registerFunctions(scriptFile, null, pigContext);
             }
         } finally {
             fis1.close();
         }

         Interpreter.setMain(true);
         FileInputStream fis = new FileInputStream(scriptFile);
         try {
             load(fis, scriptFile, pigServer.getPigContext());
         } finally {
             fis.close();
         }
         return getPigStatsMap();
     }

     /**
      * execs the script text using the interpreter. populates the pig context
      * with necessary module paths. does not close the inputstream.
      * @param script
      * @param scriptFile
      * @param pigContext
      * @throws IOException
      */
     public void load(InputStream script, String scriptFile, PigContext pigContext) throws IOException {
         Interpreter.execfile(script, scriptFile, pigContext);
     }

     @Override
     protected String getScriptingLang() {
         return "jython";
     }

     @Override
     protected Map<String, Object> getParamsFromVariables() throws IOException {
         PyFrame frame = Py.getFrame();
         @SuppressWarnings("unchecked")
         List<PyTuple> locals = ((PyStringMap) frame.getLocals()).items();
         Map<String, Object> vars = new HashMap<String, Object>();
         for (PyTuple item : locals) {
             String key = (String) item.get(0);
             Object obj = item.get(1);
             if (obj != null) {
                 String value = item.get(1).toString();
                 vars.put(key, value);
             }
         }
         return vars;
     }

     private static final Pattern p = Pattern.compile("^\\s*def\\s+(\\w+)\\s*.+");
     private static final Pattern p1 = Pattern.compile("^\\s*if\\s+__name__\\s+==\\s+[\"']__main__[\"']\\s*:\\s*$");

     private static boolean hasFunction(InputStream is) throws IOException {
         boolean hasFunction = false;
         boolean hasMain = false;
         InputStreamReader in = new InputStreamReader(is);
         BufferedReader br = new BufferedReader(in);
         String line = br.readLine();
         while (line != null) {
             if (p.matcher(line).matches()) {
                 hasFunction = true;
             } else if (p1.matcher(line).matches()) {
                 hasMain = true;
             }
             line = br.readLine();
         }
         if (hasFunction && !hasMain) {
             String msg = "Embedded script cannot mix UDFs with top level code. "
                        + "Please use if __name__ == '__main__': construct";
             throw new IOException(msg);
         }
         return hasFunction;
     }

     /**
      * File.deleteOnExit(File) does not work for a non-empty directory. This
      * Thread is used to clean up the python.cachedir (if it was a tmp dir
      * created by the Engine)
      */
     private static class DirDeleter extends Thread {
         private final File dir;
         public DirDeleter(final File file) {
             dir = file;
         }
         @Override
         public void run() {
             try {
                 delete(dir);
             } catch (Exception e) {
                 LOG.warn("on cleanup", e);
             }
         }
         private static boolean delete(final File file) {
             if (file.isDirectory()) {
                 for (File f : file.listFiles()) {
                     delete(f);
                 }
             }
             return file.delete();
         }
     }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	package org.apache.pig.scripting.jython;

	import java.io.BufferedReader;
	import java.io.File;
	import java.io.FileInputStream;
	import java.io.IOException;
	import java.io.InputStream;
	import java.io.InputStreamReader;
	import java.net.URL;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.regex.Pattern;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.pig.FuncSpec;
	import org.apache.pig.PigServer;
	import org.apache.pig.backend.executionengine.ExecException;
	import org.apache.pig.impl.PigContext;
	import org.apache.pig.impl.util.ObjectSerializer;
	import org.apache.pig.impl.util.Utils;
	import org.apache.pig.parser.ParserException;
	import org.apache.pig.scripting.ScriptEngine;
	import org.apache.pig.tools.pigstats.PigStats;
	import org.python.core.ClasspathPyImporter;
	import org.python.core.Py;
	import org.python.core.PyException;
	import org.python.core.PyFrame;
	import org.python.core.PyFunction;
	import org.python.core.PyInteger;
	import org.python.core.PyJavaPackage;
	import org.python.core.PyObject;
	import org.python.core.PyString;
	import org.python.core.PyStringMap;
	import org.python.core.PySystemState;
	import org.python.core.PyTuple;
	import org.python.modules.zipimport.zipimporter;
	import org.python.util.PythonInterpreter;

	/**
	* Implementation of the script engine for Jython
	*/
	public class JythonScriptEngine extends ScriptEngine {
	private static final Log LOG = LogFactory.getLog(JythonScriptEngine.class);

	/**
	* Language Interpreter Uses static holder pattern
	*/
	private static class Interpreter {
	static final PythonInterpreter interpreter;
	static final ArrayList<String> filesLoaded = new ArrayList<String>();
	static final String JVM_JAR;

	static {
	// should look like: file:JVM_JAR!/java/lang/Object.class
	String rpath = Object.class.getResource("Object.class").getPath();
	JVM_JAR = rpath.replaceAll("^file:(.*)!/java/lang/Object.class$", "$1");

	// Determine if a usable python.cachedir has been provided
	// if not, certain uses of jython's import will not work e.g., so create a tmp dir
	// - from some.package import *
	// - import non.jvm.package
	try {
	String skip = System.getProperty(PySystemState.PYTHON_CACHEDIR_SKIP, "false");
	if (skip.equalsIgnoreCase("true")) {
	LOG.warn("jython cachedir skipped, jython may not work");
	} else {
	File tmp = null;
	String cdir = System.getProperty(PySystemState.PYTHON_CACHEDIR);
	if (cdir != null) {
	tmp = new File(cdir);
	if (tmp.canWrite() == false) {
	LOG.error("CACHEDIR: not writable");
	throw new RuntimeException("python.cachedir not writable: " + cdir);
	}
	}
	if (tmp == null) {
	tmp = File.createTempFile("pig_jython_", "");
	tmp.delete();
	if (tmp.mkdirs() == false) {
	LOG.warn("unable to create a tmp dir for the cache, jython may not work");
	} else {
	LOG.info("created tmp python.cachedir=" + tmp);
	System.setProperty(PySystemState.PYTHON_CACHEDIR, tmp.getAbsolutePath());
	}
	Runtime.getRuntime().addShutdownHook(new DirDeleter(tmp));
	}
	}
	// local file system import path elements: current dir, JYTHON_HOME/Lib
	Py.getSystemState().path.append(new PyString(System.getProperty("user.dir")));
	String jyhome = System.getenv("JYTHON_HOME");
	if (jyhome != null) {
	Py.getSystemState().path.append(new PyString(jyhome + File.separator + "Lib"));
	}
	} catch (Exception e) {
	LOG.warn("issue with jython cache dir", e);
	}

	// cacdedir now configured, allocate the python interpreter
	interpreter = new PythonInterpreter();
	}

	/**
	* ensure the decorator functions are defined in the interpreter, and
	* manage the module import dependencies.
	* @param path location of a file to exec in the interpreter
	* @param pigContext if non-null, module import state is tracked
	* @throws IOException
	*/
	static synchronized void init(String path, PigContext pigContext) throws IOException {
	// Decorators -
	// "schemaFunction"
	// "outputSchema"
	// "outputSchemaFunction"

	if (!filesLoaded.contains(path)) {
	// attempt addition of schema decorator handler, fail silently
	interpreter.exec("def outputSchema(schema_def):\n"
	+ " def decorator(func):\n"
	+ " func.outputSchema = schema_def\n"
	+ " return func\n"
	+ " return decorator\n\n");

	interpreter.exec("def outputSchemaFunction(schema_def):\n"
	+ " def decorator(func):\n"
	+ " func.outputSchemaFunction = schema_def\n"
	+ " return func\n"
	+ " return decorator\n");

	interpreter.exec("def schemaFunction(schema_def):\n"
	+ " def decorator(func):\n"
	+ " func.schemaFunction = schema_def\n"
	+ " return func\n"
	+ " return decorator\n\n");

	InputStream is = getScriptAsStream(path);
	if (is == null) {
	throw new IllegalStateException("unable to create a stream for path: " + path);
	}
	try {
	execfile(is, path, pigContext);
	} finally {
	is.close();
	}
	}
	}

	/**
	* does not call script.close()
	* @param script
	* @param path
	* @param pigContext
	* @throws ExecException
	*/
	static void execfile(InputStream script, String path, PigContext pigContext) throws ExecException {
	try {

	if( pigContext != null ) {
	String [] argv;
	try {
	argv = (String[])ObjectSerializer.deserialize(
	pigContext.getProperties().getProperty(PigContext.PIG_CMD_ARGS_REMAINDERS));
	} catch (IOException e) {
	throw new ExecException("Cannot deserialize command line arguments", e);
	}
	PySystemState state = Py.getSystemState();
	state.argv.clear();
	if( argv != null ) {
	for (String str : argv ) {
	state.argv.append(new PyString(str));
	}
	} else {
	LOG.warn(PigContext.PIG_CMD_ARGS_REMAINDERS
	+ " is empty. This is not expected unless on testing." );
	}
	}

	// determine the current module state
	Map<String, String> before = pigContext != null ? getModuleState() : null;
	if (before != null) {
	// os.py, stax.py and posixpath.py are part of the initial state
	// if Lib directory is present and without including os.py, modules
	// which import os fail
	Set<String> includePyModules = new HashSet<String>();
	for (String key : before.keySet()) {
	// $py.class is created if Lib folder is writable
	if (key.endsWith(".py") \|\| key.endsWith("$py.class")) {
	includePyModules.add(key);
	}
	}
	before.keySet().removeAll(includePyModules);
	}

	// exec the code, arbitrary imports are processed
	interpreter.execfile(script, path);

	// determine the 'post import' module state
	Map<String, String> after = pigContext != null ? getModuleState() : null;

	// add the module files to the context
	if (after != null && pigContext != null) {
	after.keySet().removeAll(before.keySet());
	for (Map.Entry<String, String> entry : after.entrySet()) {
	String modulename = entry.getKey();
	String modulepath = entry.getValue();
	if (modulepath.equals(JVM_JAR)) {
	continue;
	} else if (modulepath.endsWith(".jar") \|\| modulepath.endsWith(".zip")) {
	pigContext.addScriptJar(modulepath);
	} else {
	pigContext.addScriptFile(modulename, modulepath);
	}
	}
	}
	} catch (PyException e) {
	if (e.match(Py.SystemExit)) {
	PyObject value = e.value;
	if (PyException.isExceptionInstance(e.value)) {
	value = value.__findattr__("code");
	}
	if (new PyInteger(0).equals(value)) {
	LOG.info("Script invoked sys.exit(0)");
	return;
	}
	}
	String message = "Python Error. " + e;
	throw new ExecException(message, 1121, e);
	}
	}

	static String get(String name) {
	return interpreter.get(name).toString();
	}

	static void setMain(boolean isMain) {
	if (isMain) {
	interpreter.set("__name__", "__main__");
	} else {
	interpreter.set("__name__", "__lib__");
	}
	}

	/**
	* get the state of modules currently loaded
	* @return a map of module name to module file (absolute path)
	*/
	private static Map<String, String> getModuleState() {
	// determine the current module state
	Map<String, String> files = new HashMap<String, String>();
	PyStringMap modules = (PyStringMap) Py.getSystemState().modules;
	for (PyObject kvp : modules.iteritems().asIterable()) {
	PyTuple tuple = (PyTuple) kvp;
	String name = tuple.get(0).toString();
	Object value = tuple.get(1);
	// inspect the module to determine file location and status
	try {
	Object fileEntry = null;
	Object loader = null;
	if (value instanceof PyJavaPackage ) {
	fileEntry = ((PyJavaPackage) value).__file__;
	} else if (value instanceof PyObject) {
	// resolved through the filesystem (or built-in)
	PyObject dict = ((PyObject) value).getDict();
	if (dict != null) {
	fileEntry = dict.__finditem__("__file__");
	loader = dict.__finditem__("__loader__");
	} // else built-in
	} // else some system module?

	if (fileEntry != null) {
	File file = resolvePyModulePath(fileEntry.toString(), loader);
	if (file.exists()) {
	String apath = file.getAbsolutePath();
	if (apath.endsWith(".jar") \|\| apath.endsWith(".zip")) {
	// jar files are simple added to the pigContext
	files.put(apath, apath);
	} else {
	// determine the relative path that the file should have in the jar
	int pos = apath.lastIndexOf(File.separatorChar + name.replace('.', File.separatorChar));
	if (pos > 0) {
	files.put(apath.substring(pos + 1), apath);
	} else {
	files.put(apath, apath);
	}
	}
	} else {
	LOG.warn("module file does not exist: " + name + ", " + file);
	}
	} // else built-in
	} catch (Exception e) {
	LOG.warn("exception while retrieving module state: " + value, e);
	}
	}
	return files;
	}
	}

	private static File resolvePyModulePath(String path, Object loader) {
	File file = new File(path);
	if (!file.exists() && loader != null) {
	if(path.startsWith(ClasspathPyImporter.PYCLASSPATH_PREFIX) && loader instanceof ClasspathPyImporter) {
	path = path.replaceFirst(ClasspathPyImporter.PYCLASSPATH_PREFIX, "");
	URL resource = ScriptEngine.class.getResource(path);
	if (resource == null) {
	resource = ScriptEngine.class.getResource(File.separator + path);
	}
	if (resource != null) {
	return new File(resource.getFile());
	}
	} else if (loader instanceof zipimporter) {
	zipimporter importer = (zipimporter) loader;
	return new File(importer.archive);
	} //JavaImporter??
	}
	return file;
	}

	@Override
	public void registerFunctions(String path, String namespace, PigContext pigContext)
	throws IOException {
	Interpreter.setMain(false);
	Interpreter.init(path, pigContext);
	pigContext.addScriptJar(getJarPath(PythonInterpreter.class));
	PythonInterpreter pi = Interpreter.interpreter;
	@SuppressWarnings("unchecked")
	List<PyTuple> locals = ((PyStringMap) pi.getLocals()).items();
	namespace = (namespace == null) ? "" : namespace + NAMESPACE_SEPARATOR;
	try {
	for (PyTuple item : locals) {
	String key = (String) item.get(0);
	Object value = item.get(1);
	FuncSpec funcspec = null;
	if (!key.startsWith("__") && !key.equals("schemaFunction")
	&& !key.equals("outputSchema")
	&& !key.equals("outputSchemaFunction")
	&& (value instanceof PyFunction)
	&& (((PyFunction)value).__findattr__("schemaFunction")== null)) {
	PyObject obj = ((PyFunction)value).__findattr__("outputSchema");
	if(obj != null) {
	Utils.getSchemaFromString(obj.toString());
	}
	funcspec = new FuncSpec(JythonFunction.class.getCanonicalName() + "('"
	+ path + "','" + key +"')");
	pigContext.registerFunction(namespace + key, funcspec);
	LOG.info("Register scripting UDF: " + namespace + key);
	}
	}
	} catch (ParserException pe) {
	throw new IOException(
	"Error parsing schema for script function from the decorator",
	pe);
	}
	pigContext.addScriptFile(path);
	Interpreter.setMain(true);
	}

	/**
	* Gets the Python function object.
	* @param path Path of the jython script file containing the function.
	* @param functionName Name of the function
	* @return a function object
	* @throws IOException
	*/
	public static PyFunction getFunction(String path, String functionName) throws IOException {
	Interpreter.setMain(false);
	Interpreter.init(path, null);
	return (PyFunction) Interpreter.interpreter.get(functionName);
	}

	@Override
	protected Map<String, List<PigStats>> main(PigContext pigContext, String scriptFile)
	throws IOException {
	if (System.getProperty(PySystemState.PYTHON_CACHEDIR_SKIP)==null)
	System.setProperty(PySystemState.PYTHON_CACHEDIR_SKIP, "false");

	PigServer pigServer = new PigServer(pigContext, false);

	// register dependencies
	String jythonJar = getJarPath(PythonInterpreter.class);
	if (jythonJar != null) {
	pigServer.registerJar(jythonJar);
	}

	File f = new File(scriptFile);

	if (!f.canRead()) {
	throw new IOException("Can't read file: " + scriptFile);
	}

	FileInputStream fis1 = new FileInputStream(scriptFile);
	try {
	if (hasFunction(fis1)) {
	registerFunctions(scriptFile, null, pigContext);
	}
	} finally {
	fis1.close();
	}

	Interpreter.setMain(true);
	FileInputStream fis = new FileInputStream(scriptFile);
	try {
	load(fis, scriptFile, pigServer.getPigContext());
	} finally {
	fis.close();
	}
	return getPigStatsMap();
	}

	/**
	* execs the script text using the interpreter. populates the pig context
	* with necessary module paths. does not close the inputstream.
	* @param script
	* @param scriptFile
	* @param pigContext
	* @throws IOException
	*/
	public void load(InputStream script, String scriptFile, PigContext pigContext) throws IOException {
	Interpreter.execfile(script, scriptFile, pigContext);
	}

	@Override
	protected String getScriptingLang() {
	return "jython";
	}

	@Override
	protected Map<String, Object> getParamsFromVariables() throws IOException {
	PyFrame frame = Py.getFrame();
	@SuppressWarnings("unchecked")
	List<PyTuple> locals = ((PyStringMap) frame.getLocals()).items();
	Map<String, Object> vars = new HashMap<String, Object>();
	for (PyTuple item : locals) {
	String key = (String) item.get(0);
	Object obj = item.get(1);
	if (obj != null) {
	String value = item.get(1).toString();
	vars.put(key, value);
	}
	}
	return vars;
	}

	private static final Pattern p = Pattern.compile("^\\sdef\\s+(\\w+)\\s.+");
	private static final Pattern p1 = Pattern.compile("^\\sif\\s+__name__\\s+==\\s+[\"']__main__[\"']\\s:\\s*$");

	private static boolean hasFunction(InputStream is) throws IOException {
	boolean hasFunction = false;
	boolean hasMain = false;
	InputStreamReader in = new InputStreamReader(is);
	BufferedReader br = new BufferedReader(in);
	String line = br.readLine();
	while (line != null) {
	if (p.matcher(line).matches()) {
	hasFunction = true;
	} else if (p1.matcher(line).matches()) {
	hasMain = true;
	}
	line = br.readLine();
	}
	if (hasFunction && !hasMain) {
	String msg = "Embedded script cannot mix UDFs with top level code. "
	+ "Please use if __name__ == '__main__': construct";
	throw new IOException(msg);
	}
	return hasFunction;
	}

	/**
	* File.deleteOnExit(File) does not work for a non-empty directory. This
	* Thread is used to clean up the python.cachedir (if it was a tmp dir
	* created by the Engine)
	*/
	private static class DirDeleter extends Thread {
	private final File dir;
	public DirDeleter(final File file) {
	dir = file;
	}
	@Override
	public void run() {
	try {
	delete(dir);
	} catch (Exception e) {
	LOG.warn("on cleanup", e);
	}
	}
	private static boolean delete(final File file) {
	if (file.isDirectory()) {
	for (File f : file.listFiles()) {
	delete(f);
	}
	}
	return file.delete();
	}
	}
	}