lil' python obj plucking util

from_obj_path.py

import unittest
from typing import List

FALLBACK_FLAG = '_______FALLBACK_FLAG'

def from_obj_path(obj, path, fallback = None):
  """ from_obj_path({'name': {'first': 'bo'}}, 'name.first') -> 'bo'"""
  d = dict(obj)
  if '.' not in path:
    val = d.get(path, FALLBACK_FLAG)
    if val == FALLBACK_FLAG:
      return fallback, Exception(path + ' does not exist in {}'.format(d))
    return val, None
  paths = path.split('.')
  first, rest = paths[0], paths[1:]
  val = d.get(first, FALLBACK_FLAG)
  if val == FALLBACK_FLAG:
    return fallback, Exception(first + ' does not exist in {}'.format(d))
  if not isinstance(val, dict):
    return val, Exception('should be dict')
  return from_obj_path(val, '.'.join(rest), fallback)

class TestPluck(unittest.TestCase):
  def test_simple(self):
    arg = {'a': 1}
    result, _err = from_obj_path(arg, 'a')
    expectation = 1
    self.assertEqual(result, expectation)

  def test_not_found(self):
    arg = {'a': 1}
    _result, err = from_obj_path(arg, 'b')
    self.assertRegex(str(err) ,r'does not exist')

  def test_nested(self):
    arg = {'a': {'a1': 1}}
    result, _err = from_obj_path(arg, 'a.a1')
    self.assertEqual(result, 1)
    arg = {'a': {'a1': {'a2': 2}}}
    result, _err = from_obj_path(arg, 'a.a1.a2')
    self.assertEqual(result, 2)
  
  def test_fallback(self):
    arg = {'a': {'a1': 1}}
    fallback = 'FALLBACK'
    result, _err = from_obj_path(arg, 'a.a2', fallback)
    self.assertEqual(result, fallback)

unittest.main()