recipe/promise_plugin: Improve options

slapos/recipe/promise_plugin.py

@@ -24,15 +24,17 @@
 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 #
 ##############################################################################
+import ast
 import re
 import logging, os
 import zc.buildout.easy_install
 from pprint import pformat
 from slapos.recipe.librecipe import GenericBaseRecipe
+from zc.buildout import UserError
 
 script_template = '''# This script is auto generated by slapgrid, do not edit!
 import sys
-sys.path[0:0] = %(path)s
+sys.path[0:0] = %(paths)s
 
 extra_config_dict = %(config)s
 
@@ -54,6 +56,24 @@ for module in list(sys.modules):
 %(content)s
 '''
 
+execfile_template = """\
+def _(path):
+  global _
+  del _
+  with open(path) as f:
+    return compile(f.read(), path, 'exec')
+exec(_(%r))"""
+
+
+def isPythonName(name):
+  try:
+    x, = ast.parse(name).body
+    if type(x) is ast.Expr:
+      x = x.value
+      return type(x) is ast.Name and x.id == name
+  except Exception:
+    pass
+
 class Recipe(GenericBaseRecipe):
 
   _WORKING_SET_CACHE_NAME = "slapos.cookbook_pplugin_ws_cache"
@@ -64,6 +84,35 @@ class Recipe(GenericBaseRecipe):
     options['develop-eggs-directory'] = buildout_section['develop-eggs-directory']
     super(Recipe, self).__init__(buildout, name, options)
 
+    filepath = options.get('file')
+    module = options.get('module')
+    klass = options.get('class')
+
+    if klass == 'RunPromise':
+      klass = None
+    elif klass and not isPythonName(klass):
+      raise UserError("%r is not valid class name" % klass)
+
+    if bool(module) == bool(filepath):
+      raise UserError("Either 'module' or 'file' is required but not both")
+
+    if module:
+      if not all(map(isPythonName, module.split('.'))):
+        raise UserError("%r is not a valid module name" % module)
+      if klass:
+        content = 'from %s import %s as RunPromise' % (module, klass)
+      else:
+        content = 'from %s import RunPromise' % module
+    else:
+      content = execfile_template % filepath
+      if klass:
+        content += '\n\nRunPromise = ' + klass
+    self.content = content
+
+    self.config_dict = {key[7:]: options[key]
+      for key in options
+      if key.startswith('config-')}
+
   def _get_cache_storage(self):
     """Return a mapping where to store generated working sets.
     from https://github.com/buildout/buildout/blob/master/zc.recipe.egg_/src/zc/recipe/egg/egg.py#L170
@@ -111,26 +160,10 @@ class Recipe(GenericBaseRecipe):
       else:
         working_set = set()
 
-    regex = r"^[\w_\-\.\s]+$"
-    import_path = self.options.get('import', '').strip()
-    if import_path:
-      if not re.search(regex, import_path):
-        raise ValueError("Import path %r is not a valid" % import_path)
-      content_string = "from %s import RunPromise" % import_path
-    else:
-      # old parameter for compatibility
-      content_string = self.options['content'].strip()
-      if not re.search(regex, content_string):
-        raise ValueError("Promise content %r is not valid" % content_string)
-
-    config_dict = {key[7:]: self.options[key]
-      for key in self.options
-      if key.startswith('config-')}
-
     return self.createFile(self.options['output'], script_template % {
-      'path': pformat([dist.location for dist in working_set], indent=2),
-      'content': content_string,
-      'config': pformat(config_dict, indent=2),
+      'paths': pformat(tuple(dist.location for dist in working_set), indent=2),
+      'config': pformat(self.config_dict, indent=2),
+      'content': self.content,
       }, int(self.options.get('mode', '0644'), 8)),
 
   update = install