key.py

import copy
from itertools import izip

from gcloud.datastore import datastore_v1_pb2 as datastore_pb
from gcloud.datastore.dataset import Dataset


class Key(object):
  """
  An immutable representation of a datastore Key.
  """

  def __init__(self, dataset=None, namespace=None, path=None):
    self._dataset = dataset
    self._namespace = namespace
    self._path = path or [{'kind': ''}]

  def _clone(self):
    """Duplicates the Key.

    We make a shallow copy of the :class:`gcloud.datastore.dataset.Dataset`
    because it holds a reference an authenticated connection,
    which we don't want to lose.
    """
    clone = copy.deepcopy(self)
    clone._dataset = self._dataset  # Make a shallow copy of the Dataset.
    return clone

  @classmethod
  def from_protobuf(cls, pb, dataset=None):
    path = []
    for element in pb.path_element:
      element_dict = {'kind': element.kind}

      if element.HasField('id'):
        element_dict['id'] = element.id

      elif element.HasField('name'):
        element_dict['name'] = element.name

      path.append(element_dict)

    if not dataset:
      dataset = Dataset(id=pb.partition_id.dataset_id)

    return cls(path=path, dataset=dataset)

  def to_protobuf(self):
    key = datastore_pb.Key()

    # Technically a dataset is required to do anything with the key,
    # but we shouldn't throw a cryptic error if one isn't provided
    # in the initializer.
    if self.dataset():
      # Apparently 's~' is a prefix for High-Replication and is necessary here.
      # Another valid preflix is 'e~' indicating EU datacenters.
      dataset_id = self.dataset().id()
      if dataset_id:
        if dataset_id[:2] not in ['s~', 'e~']:
          dataset_id = 's~' + dataset_id

        key.partition_id.dataset_id = dataset_id

    if self._namespace:
      key.partition_id.namespace = self._namespace

    for item in self.path():
      element = key.path_element.add()
      if 'kind' in item:
        element.kind = item['kind']
      if 'id' in item:
        element.id = item['id']
      if 'name' in item:
        element.name = item['name']

    return key

  @classmethod
  def from_path(cls, *args, **kwargs):
    path = []
    items = iter(args)

    for kind, id_or_name in izip(items, items):
      entry = {'kind': kind}
      if isinstance(id_or_name, basestring):
        entry['name'] = id_or_name
      else:
        entry['id'] = id_or_name
      path.append(entry)

    kwargs['path'] = path
    return cls(**kwargs)

  def is_partial(self):
    return (self.id_or_name() is None)

  def dataset(self, dataset=None):
    if dataset:
      clone = self._clone()
      clone._dataset = dataset
      return clone
    else:
      return self._dataset

  def namespace(self, namespace=None):
    if namespace:
      clone = self._clone()
      clone._namespace = namespace
      return clone
    else:
      return self._namespace

  def path(self, path=None):
    if path:
      clone = self._clone()
      clone._path = path
      return clone
    else:
      return self._path

  def kind(self, kind=None):
    if kind:
      clone = self._clone()
      clone._path[-1]['kind'] = kind
      return clone
    elif self.path():
      return self._path[-1]['kind']

  def id(self, id=None):
    if id:
      clone = self._clone()
      clone._path[-1]['id'] = id
      return clone
    elif self.path():
      return self._path[-1].get('id')

  def name(self, name=None):
    if name:
      clone = self._clone()
      clone._path[-1]['name'] = name
      return clone
    elif self.path():
      return self._path[-1].get('name')

  def id_or_name(self):
    return self.id() or self.name()

  def parent(self):#pragma NO COVER
    # See https://github.com/GoogleCloudPlatform/gcloud-python/issues/135
    raise NotImplementedError

  def __repr__(self): #pragma NO COVER
    return '<Key%s>' % self.path()