Add $ref support to policy schema

components/policy/tools/generate_policy_source.py

 #!/usr/bin/env python
 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 '''python %prog [options] platform chromium_os_flag template
 
 platform specifies which platform source is being generated for
   and can be one of (win, mac, linux)
 chromium_os_flag should be 1 if this is a Chromium OS build
 template is the path to a .json policy template file.'''
 
 from __future__ import with_statement
+from functools import partial
 import json
 from optparse import OptionParser
 import re
 import sys
 import textwrap
+import types
 
 
 CHROME_POLICY_KEY = 'SOFTWARE\\\\Policies\\\\Google\\\\Chrome'
 CHROMIUM_POLICY_KEY = 'SOFTWARE\\\\Policies\\\\Chromium'
 
 
 class PolicyDetails:
   """Parses a policy template and caches all its details."""
 
   # Maps policy types to a tuple with 3 other types:
@@ skipping 246 matching lines @@
     self.string_enums = []
     self.simple_types = {
       'boolean': None,
       'integer': None,
       'null': None,
       'number': None,
       'string': None,
     }
     self.stringlist_type = None
     self.ranges = {}
+    self.id_map = {}
 
   def GetString(self, s):
     if s in self.shared_strings:
       return self.shared_strings[s]
     # Generate JSON escaped string, which is slightly different from desired
     # C/C++ escaped string. Known differences includes unicode escaping format.
     return json.dumps(s)
 
   def AppendSchema(self, type, extra, comment=''):
     index = len(self.schema_nodes)
@@ skipping 96 matching lines @@
         str(min_value) if min_value_set else 'INT_MIN')
     return self.AppendSchema('TYPE_INTEGER',
         index,
         'integer with ranged restriction: %s' % name)
 
   def Generate(self, schema, name):
     """Generates the structs for the given schema.
 
     |schema|: a valid JSON schema in a dictionary.
     |name|: the name of the current node, for the generated comments."""
+    if schema.has_key('$ref'):
+      if schema.has_key('id'):
+        raise RuntimeError("Schemas with a $ref can't have an id")
+      if not isinstance(schema['$ref'], types.StringTypes):
+        raise RuntimeError("$ref attribute must be a string")
+      return schema['$ref']
     if schema['type'] in self.simple_types:
       if not self.SchemaHaveRestriction(schema):
         # Simple types use shared nodes.
         return self.GetSimpleType(schema['type'])
       elif 'enum' in schema:
         return self.GetEnumType(schema, name)
       elif 'pattern' in schema:
         return self.GetPatternType(schema, name)
       else:
         return self.GetRangedType(schema, name)
 
     if schema['type'] == 'array':
       # Special case for lists of strings, which is a common policy type.
       if schema['items']['type'] == 'string':
         return self.GetStringList()
-      return self.AppendSchema(
-          'TYPE_LIST',
-          self.Generate(schema['items'], 'items of ' + name))
+      return self.AppendSchema('TYPE_LIST',
+          self.GenerateAndCollectID(schema['items'], 'items of ' + name))
     elif schema['type'] == 'object':
       # Reserve an index first, so that dictionaries come before their
       # properties. This makes sure that the root node is the first in the
       # SchemaNodes array.
       index = self.AppendSchema('TYPE_DICTIONARY', -1)
 
       if 'additionalProperties' in schema:
-        additionalProperties = self.Generate(
+        additionalProperties = self.GenerateAndCollectID(
             schema['additionalProperties'],
             'additionalProperties of ' + name)
       else:
         additionalProperties = -1
 
       # Properties must be sorted by name, for the binary search lookup.
       # Note that |properties| must be evaluated immediately, so that all the
       # recursive calls to Generate() append the necessary child nodes; if
       # |properties| were a generator then this wouldn't work.
       sorted_properties = sorted(schema.get('properties', {}).items())
-      properties = [ (self.GetString(key), self.Generate(subschema, key))
-                     for key, subschema in sorted_properties ]
+      properties = [
+          (self.GetString(key), self.GenerateAndCollectID(subschema, key))
+          for key, subschema in sorted_properties ]
 
       pattern_properties = []
       for pattern, subschema in schema.get('patternProperties', {}).items():
         pattern_properties.append((self.GetString(pattern),
-            self.Generate(subschema, pattern)));
+            self.GenerateAndCollectID(subschema, pattern)));
 
       begin = len(self.property_nodes)
       self.property_nodes += properties
       end = len(self.property_nodes)
       self.property_nodes += pattern_properties
       pattern_end = len(self.property_nodes)
 
       if index == 0:
         self.root_properties_begin = begin
         self.root_properties_end = end
 
       extra = len(self.properties_nodes)
       self.properties_nodes.append((begin, end, pattern_end,
           additionalProperties, name))
 
       # Set the right data at |index| now.
       self.schema_nodes[index] = ('TYPE_DICTIONARY', extra, name)
       return index
     else:
       assert False
 
+  def GenerateAndCollectID(self, schema, name):
+    """A wrapper of Generate(), will take the return value, check and add 'id'
+    attribute to self.id_map. The wrapper needs to be used for every call to
+    Generate().
+    """
+    index = self.Generate(schema, name)
+    if not schema.has_key('id'):
+      return index
+    id_str = schema['id']
+    if self.id_map.has_key(id_str):
+      raise RuntimeError('Duplicated id: ' + id_str)
+    self.id_map[id_str] = index
+    return index
+
   def Write(self, f):
     """Writes the generated structs to the given file.
 
     |f| an open file to write to."""
     f.write('const internal::SchemaNode kSchemas[] = {\n'
             '//  Type                          Extra\n')
     for type, extra, comment in self.schema_nodes:
       type += ','
       f.write('  { base::Value::%-18s %3d },  // %s\n' % (type, extra, comment))
     f.write('};\n\n')
@@ skipping 34 matching lines @@
 
     f.write('const internal::SchemaData kChromeSchemaData = {\n'
             '  kSchemas,\n')
     f.write('  kPropertyNodes,\n' if self.property_nodes else '  NULL,\n')
     f.write('  kProperties,\n' if self.properties_nodes else '  NULL,\n')
     f.write('  kRestrictionNodes,\n' if self.restriction_nodes else '  NULL,\n')
     f.write('  kIntegerEnumerations,\n' if self.int_enums else '  NULL,\n')
     f.write('  kStringEnumerations,\n' if self.string_enums else '  NULL,\n')
     f.write('};\n\n')
 
+  def GetByID(self, id_str):
+    if not isinstance(id_str, types.StringTypes):
+      return id_str
+    if not self.id_map.has_key(id_str):
+      raise RuntimeError('Invalid $ref: ' + id_str)
+    return self.id_map[id_str]
+
+  def ResolveID(self, index, params):
+    return params[:index] + (self.GetByID(params[index]),) + params[index+1:]
+
+  def ResolveReferences(self):
+    """Resolve reference mapping, required to be called after Generate()
+
+    After calling Generate(), the type of indices used in schema structures
+    might be either int or string. An int type suggests that it's a resolved
+    index, but for string type it's unresolved. Resolving a reference is as
+    simple as looking up for corresponding ID in self.id_map, and replace the
+    old index with the mapped index.
+    """
+    self.schema_nodes = map(partial(self.ResolveID, 1), self.schema_nodes)
+    self.property_nodes = map(partial(self.ResolveID, 1), self.property_nodes)
+    self.properties_nodes = map(partial(self.ResolveID, 3),
+        self.properties_nodes)
 
 def _WritePolicyConstantSource(policies, os, f):
   f.write('#include "policy/policy_constants.h"\n'
           '\n'
           '#include <algorithm>\n'
           '#include <climits>\n'
           '\n'
           '#include "base/logging.h"\n'
           '#include "components/policy/core/common/schema_internal.h"\n'
           '\n'
@@ skipping 22 matching lines @@
   for policy in policies:
     if policy.is_supported:
       f.write('  { %-14s %-16s %3s, %24s },\n' % (
                   'true,' if policy.is_deprecated else 'false,',
                   'true,' if policy.is_device_only else 'false,',
                   policy.id,
                   policy.max_size))
   f.write('};\n\n')
 
   schema_generator = SchemaNodesGenerator(shared_strings)
-  schema_generator.Generate(chrome_schema, 'root node')
+  schema_generator.GenerateAndCollectID(chrome_schema, 'root node')
+  schema_generator.ResolveReferences()
   schema_generator.Write(f)
 
   f.write('bool CompareKeys(const internal::PropertyNode& node,\n'
           '                 const std::string& key) {\n'
           '  return node.key < key;\n'
           '}\n\n')
 
   f.write('}  // namespace\n\n')
 
   if os == 'win':
@@ skipping 298 matching lines @@
 def _WriteCloudPolicyDecoder(policies, os, f):
   f.write(CPP_HEAD)
   for policy in policies:
     if policy.is_supported and not policy.is_device_only:
       _WritePolicyCode(f, policy)
   f.write(CPP_FOOT)
 
 
 if __name__ == '__main__':
   sys.exit(main())