Query the type of a keypath in a JSON Schema

Question

This was the question given to me:

Implement a function to query json_schema

The goal of this exercise is to implement a function that allows you to query the type of a keypath in a JSON schema.

This function will accept a valid JSON schema as dict, a key_path (eg: foo.bar.baz) and return the type of the property.

Note:

• There are only two fields in the schema you have to pay attention to: properties and definitions.
• If the dictionary associated with the field has a field named $ref it means that it's referring to another schema stored under the top- level schema. You have follow the link to get to the actual definition.
• For the sake of this exercise you can assume that all values for $ref will start with #/<key_path>.
• You should see a schema and some assert statements under "Test" section. You should NOT have to change anything under the "Test" section. If you can get our code to pass the tests, then it means your function works as expected.
• Feel free to use whatever libraries you need to use except any libraries that actually allow you to query the JSON schema.

I followed it by doing this:

import json
import copy

schema = json.loads('''{
  "$id": "https://example.com/nested-schema.json",
  "title": "nested-schema",
  "$schema": "http://json-schema.org/draft-07/schema#",
  "required": [
    "EmploymentInformation",
    "EmployeePartyID",
    "Age"
  ],
  "properties": {
    "EmployeePartyID": {
      "type": "string",
      "minLength": 1,
      "maxLength": 3
    },
    "EmploymentInformation": {
      "$ref": "#/definitions/EmploymentInformation"
    },
    "Age": {
      "type": "integer",
      "minimum": 16,
      "maximum": 80
    }
  },
  "definitions": {
    "EmploymentInformation": {
      "type": "object",
      "required": [
        "OriginalHireDate"
      ],
      "properties": {
        "OriginalHireDate": {
          "type": "string",
          "format": "date"
        },
        "Beneficiary": {
          "$ref": "#/definitions/DependantInformation"
        }
      }
    },
    "DependantInformation": {
      "type": "object",
      "required": [
        "Name"
      ],
      "properties": {
        "Name": {
          "type": "string",
          "minLength": 5
        }
      }
    }
  },
  "description": "nested-schema"
}''')


def resolve_ref(ref, modified_schema):
    ref_path = ref["$ref"].split("/")[1:]
    ref_obj = modified_schema
    for node in ref_path:
        ref_obj = ref_obj[node]
    resolve_refs(ref_obj, modified_schema)
    return ref_obj


def resolve_refs(json_schema, modified_schema=None):
    if modified_schema is None:
        modified_schema = json_schema

    for k, v in json_schema.items():
        if isinstance(v, dict) and "$ref" in v:
            json_schema[k] = resolve_ref(v, modified_schema)
        elif isinstance(v, dict):
            resolve_refs(json_schema[k], modified_schema)


def get_type(key_path, json_schema):
    """
    Recursively gets the type if it exists.
    :param key_path:
    :param json_schema:
    :return:
    """
    if 'properties' in json_schema:
        if key_path[0] in json_schema['properties']:
            return get_type(key_path[1:], json_schema['properties'][key_path[0]])
        else:
            return json_schema.get('type', None)
    else:
        return json_schema.get('type', None)


completed_schema = dict()
def get_complete_schema(json_schema):
    """
    Takes the schema and solves the refs, stores in a dict so that it is computed only once.
    :param json_schema:
    :return modified schema:
    """
    schema_str = json.dumps(json_schema, sort_keys=True)
    if not completed_schema.get(schema_str):
        modified_schema = copy.deepcopy(json_schema)
        resolve_refs(modified_schema)
        completed_schema[schema_str] = modified_schema
    else:
        modified_schema = completed_schema[schema_str]
    return modified_schema


def get_type_for_key_path(json_schema: dict, key_path: str) -> str:
    modified_schema = get_complete_schema(json_schema)
    key_path_list = key_path.split('.')
    key_path_type = get_type(key_path_list, modified_schema)
    return key_path_type


assert (get_type_for_key_path(schema, "Age") == "integer")
assert (get_type_for_key_path(schema, "EmploymentInformation.OriginalHireDate") == "string")
assert (get_type_for_key_path(schema, "EmploymentInformation.Beneficiary.Name") == "string")
assert (get_type_for_key_path(schema, "foo.bar") == None)

Would my code be good enough for a Senior Python Engineer?

Note, they gave me the schema and the test conditions.

Answer 1

• You have failed to apply TDD, and have over complicated the solution.
  Before reviewing your code I completed the challenge myself, which I'll walk you through after reviewing your solution.

• Your function resolve_ref

  • This should be split out into two functions, get_ref and _resolve_ref, having resolve_ref as a convenience to them.
  • Since you haven't split out get_ref, you have to mangle the object to not contain $ref in the walk part of the function.
  • I'm not a fan of defining ref_path, just place that next to the in in the for loop.
  • You shouldn't call resolve_refs in resolve_ref it's making you walk the dictionary tons of times, and it's just plain confusing.

• Your function resolve_refs

  • It's merged _resolve_ref into it making it far more complex than it needs to be. I suggest splitting this out.

• Your function get_type

  • If I were marking your code and saw this function, you'd be dropped within seconds.
  • If I'm asking for an object's type, I'm not asking for its parent's type.
  • This should be split into two functions, get_property and get_type. The former walks the path getting the property, and the latter just uses get_property(...).get('type', None).

• Your function get_complete_schema

  • Never mind an interview, if I were a maintainer for a JSONSchema library, and you submitted this in a pull request. I'd never accept it, whilst this function exists.
  • You rely on a global completed_schema which means your function works once, and then blows up every time after that.
  • You've been asked to get a property, not make a new schema.
  • If we need to edit the schema after you've had you way with it, then we now have to keep track of all the references and provide an annoying interface to modify the schema. It also leads to data integrity problems if a user or the library provider messes up once.

• Your function get_type_for_key_path

  • This seems reasonable enough, however key_path.split('.') should probably be in get_type.

• Your naming sense is poor, I don't want to read 8 character variable names when a 4 character one is enough.

• You've quarter arsed your docstrings, and they're not even PEP compliant. If this is the level of documentation you'd give when you provide documentation then I'd not want you. You've given a short description on what it does, and stated what parameters it takes, but not explained the parameters. Additionally only half your code is documented.

  If you're going to do something, at least do it well.

How I solved this

1. Get the code working with Age.

   1. Make the get_type_from_key_path.

      Given the god awful name, we know it's a convenience function. This means we should delegate to a different function to get the property, and this should only mutate the result to pass the tests.

   2. Make get_property.

      This splits the provided path into segments, and walks the tree. It should be noted that each time you walk here you're walking node['properties'][segment] not node[segment].

   def get_property(schema, path):
       node = schema
       for segment in path.split('.'):
           node = node['properties'][segment]
       return node
   
   
   def get_type_from_key_path(schema, path):
       return get_property(schema, path)['type']
   

2. Get the code working with EmploymentInformation.OriginalHireDate.

   1. We need to add a new function to resolve references. Since this is a programming challenge we can look to the problem description to make things simple.

      For the sake of this exercise you can assume that all values for $ref will start with #/.

      This means we only need to pass the schema and walk the provided path.

   2. Change get_property so if "$ref" is in the object to change the node to the reference.

   def get_ref(schema, path):
       node = schema
       for segment in path.split('/')[1:]:
           node = node[segment]
       return node
   
   
   def get_property(schema, path):
       node = schema
       for segment in path.split('.'):
           if '$ref' in node:
               node = get_ref(schema, node['$ref'])
           node = node['properties'][segment]
       return node
   

3. Get the code working with EmploymentInformation.Beneficiary.Name.
   No changes needed, it just works!

4. Get the code working with foo.bar.

   1. Change the code so if get_property raises a key error then you return None.

   def get_type_for_key_path(schema, path):
       try:
           return get_property(schema, path)['type']
       except KeyError:
           return None
   

This nicely works with all the tests and is really short.

def get_ref(schema, path):
    node = schema
    for segment in path.split('/')[1:]:
        node = node[segment]
    return node


def get_property(schema, path):
    node = schema
    for segment in path.split('.'):
        if '$ref' in node:
            node = get_ref(schema, node['$ref'])
        node = node['properties'][segment]
    return node


def get_type_for_key_path(schema, path):
    try:
        return get_property(schema, path)['type']
    except KeyError:
        return None