<commit_before>import mimetypes
import os
import re
from datetime import datetime
from decimal import Decimal

from django.utils.timezone import utc

from photonix.photos.models import Camera, Lens, Photo, PhotoFile, Task, Library, Tag, PhotoTag
from photonix.photos.utils.metadata import (PhotoMetadata, parse_datetime, parse_gps_location)

def record_photo(path, library, inotify_event_type=None):
    if type(library) == Library:
        library_id = library.id
    else:
        library_id = str(library)
    try:
        photo_file = PhotoFile.objects.get(path=path)
    except PhotoFile.DoesNotExist:
        photo_file = PhotoFile()

    if inotify_event_type in ['DELETE', 'MOVED_FROM']:
        if PhotoFile.objects.filter(path=path).exists():
            return delete_photo_record(photo_file)
        else:
            return True

    file_modified_at = datetime.fromtimestamp(os.stat(path).st_mtime, tz=utc)

    if photo_file and photo_file.file_modified_at == file_modified_at:
        return True

    metadata = PhotoMetadata(path)
    date_taken = None
    possible_date_keys = ['Date/Time Original', 'Date Time Original', 'Date/Time', 'Date Time', 'GPS Date/Time', 'Modify Date', 'File Modification Date/Time']
    for date_key in possible_date_keys:
        date_taken = parse_datetime(metadata.get(date_key))
        if date_taken:
            break

    camera = None
    camera_make = metadata.get('Make', '')
    camera_model = metadata.get('Camera Model Name', '')
    if camera_model:
        camera_model = camera_model.replace(camera_make, '').strip()
    if camera_make and camera_model:
        try:
            camera = Camera.objects.get(library_id=library_id, make=camera_make, model=camera_model)
            if date_taken < camera.earliest_photo:
                camera.earliest_photo = date_taken
                camera.save()
            if date_taken > camera.latest_photo:
                camera.latest_photo = date_taken
                camera.save()
        except Camera.DoesNotExist:
            camera = Camera(library_id=library_id, make=camera_make, model=camera_model,
                            earliest_photo=date_taken, latest_photo=date_taken)
            camera.save()

    lens = None
    lens_name = metadata.get('Lens ID')
    if lens_name:
        try:
            lens = Lens.objects.get(name=lens_name)
            if date_taken < lens.earliest_photo:
                lens.earliest_photo = date_taken
                lens.save()
            if date_taken > lens.latest_photo:
                lens.latest_photo = date_taken
                lens.save()
        except Lens.DoesNotExist:
            lens = Lens(library_id=library_id, name=lens_name, earliest_photo=date_taken,
                        latest_photo=date_taken)
            lens.save()

    photo = None
    if date_taken:
        try:
            # TODO: Match on file number/file name as well
            photo = Photo.objects.get(taken_at=date_taken)
        except Photo.DoesNotExist:
            pass

    latitude = None
    longitude = None
    if metadata.get('GPS Position'):
        latitude, longitude = parse_gps_location(metadata.get('GPS Position'))

    iso_speed = None
    if metadata.get('ISO'):
        try:
            iso_speed = int(re.search(r'[0-9]+', metadata.get('ISO')).group(0))
        except AttributeError:
            pass
    if not photo:
        # Save Photo
        aperture = None
        aperturestr = metadata.get('Aperture')
        if aperturestr:
            try:
                aperture = Decimal(aperturestr)
                if aperture.is_infinite():
                    aperture = None
            except:
                pass

        photo = Photo(
            library_id=library_id,
            taken_at=date_taken,
            taken_by=metadata.get('Artist') or None,
            aperture=aperture,
            exposure=metadata.get('Exposure Time') or None,
            iso_speed=iso_speed,
            focal_length=metadata.get('Focal Length') and metadata.get('Focal Length').split(' ', 1)[0] or None,
            flash=metadata.get('Flash') and 'on' in metadata.get('Flash').lower() or False,
            metering_mode=metadata.get('Metering Mode') or None,
            drive_mode=metadata.get('Drive Mode') or None,
            shooting_mode=metadata.get('Shooting Mode') or None,
            camera=camera,
            lens=lens,
            latitude=latitude,
            longitude=longitude,
            altitude=metadata.get('GPS Altitude') and metadata.get('GPS Altitude').split(' ')[0],
            star_rating=metadata.get('Rating') or 0    # Created at 27-4-2021: To fetch tag from metadata and assign to star rating.
        )
        photo.save()

        # Created at 27-4-2021: To fetch tags from metadata and create tag for photo.
        subject = metadata.get('Subject')
        if subject:
            tag = Tag.objects.create(library_id=library_id, name=subject, type="G")
            PhotoTag.objects.create(
                photo=photo,
                tag=tag,
                confidence=1.0
            )
    else:
        for photo_file in photo.files.all():
            if not os.path.exists(photo_file.path):
                photo_file.delete()

    width = metadata.get('Image Width')
    height = metadata.get('Image Height')
    if metadata.get('Orientation') in ['Rotate 90 CW', 'Rotate 270 CCW', 'Rotate 90 CCW', 'Rotate 270 CW']:
        old_width = width
        width = height
        height = old_width

    # Save PhotoFile
    photo_file.photo = photo
    photo_file.path = path
    photo_file.width = width
    photo_file.height = height
    photo_file.mimetype = mimetypes.guess_type(path)[0]
    photo_file.file_modified_at = file_modified_at
    photo_file.bytes = os.stat(path).st_size
    photo_file.preferred = False  # TODO
    photo_file.save()

    # Create task to ensure JPEG version of file exists (used for thumbnailing, analysing etc.)
    Task(
        type='ensure_raw_processed',
        subject_id=photo.id,
        complete_with_children=True
    ).save()

    return photo

def delete_photo_record(photo_file_obj):
    """Delete photo record if photo not exixts on library path."""
    delete_photofile_and_photo_record(photo_file_obj)
    Tag.objects.filter(photo_tags=None).delete()
    Camera.objects.filter(photos=None).delete()
    Lens.objects.filter(photos=None).delete()
    return True

def move_or_rename_photo(photo_old_path, photo_new_path, library_id):
    """Rename a photoFile or change the path while moving photo in child directory."""
    try:
        photo_file = PhotoFile.objects.get(path=photo_old_path)
        photo_file.path = photo_new_path
        photo_file.save()
        return photo_file
    except Exception as e:
        return True

def delete_child_dir_all_photos(directory_path, library_id):
    """When a child directory deleted it delete all the photo records of that directory."""
    for photo_file_obj in PhotoFile.objects.filter(path__startswith=directory_path):
        delete_photofile_and_photo_record(photo_file_obj)
    Tag.objects.filter(photo_tags=None).delete()
    Camera.objects.filter(photos=None).delete()
    Lens.objects.filter(photos=None).delete()
    return True

def delete_photofile_and_photo_record(photo_file_obj):
    """Delete photoFile object with its photo object."""
    photo_obj = photo_file_obj.photo
    photo_file_obj.delete()
    if not photo_obj.files.all():
        photo_obj.delete()
<commit_msg>Separate multiple subject tags by comma<commit_after>import mimetypes
import os
import re
from datetime import datetime
from decimal import Decimal

from django.utils.timezone import utc

from photonix.photos.models import Camera, Lens, Photo, PhotoFile, Task, Library, Tag, PhotoTag
from photonix.photos.utils.metadata import (PhotoMetadata, parse_datetime, parse_gps_location)

def record_photo(path, library, inotify_event_type=None):
    if type(library) == Library:
        library_id = library.id
    else:
        library_id = str(library)
    try:
        photo_file = PhotoFile.objects.get(path=path)
    except PhotoFile.DoesNotExist:
        photo_file = PhotoFile()

    if inotify_event_type in ['DELETE', 'MOVED_FROM']:
        if PhotoFile.objects.filter(path=path).exists():
            return delete_photo_record(photo_file)
        else:
            return True

    file_modified_at = datetime.fromtimestamp(os.stat(path).st_mtime, tz=utc)

    if photo_file and photo_file.file_modified_at == file_modified_at:
        return True

    metadata = PhotoMetadata(path)
    date_taken = None
    possible_date_keys = ['Date/Time Original', 'Date Time Original', 'Date/Time', 'Date Time', 'GPS Date/Time', 'Modify Date', 'File Modification Date/Time']
    for date_key in possible_date_keys:
        date_taken = parse_datetime(metadata.get(date_key))
        if date_taken:
            break

    camera = None
    camera_make = metadata.get('Make', '')
    camera_model = metadata.get('Camera Model Name', '')
    if camera_model:
        camera_model = camera_model.replace(camera_make, '').strip()
    if camera_make and camera_model:
        try:
            camera = Camera.objects.get(library_id=library_id, make=camera_make, model=camera_model)
            if date_taken < camera.earliest_photo:
                camera.earliest_photo = date_taken
                camera.save()
            if date_taken > camera.latest_photo:
                camera.latest_photo = date_taken
                camera.save()
        except Camera.DoesNotExist:
            camera = Camera(library_id=library_id, make=camera_make, model=camera_model,
                            earliest_photo=date_taken, latest_photo=date_taken)
            camera.save()

    lens = None
    lens_name = metadata.get('Lens ID')
    if lens_name:
        try:
            lens = Lens.objects.get(name=lens_name)
            if date_taken < lens.earliest_photo:
                lens.earliest_photo = date_taken
                lens.save()
            if date_taken > lens.latest_photo:
                lens.latest_photo = date_taken
                lens.save()
        except Lens.DoesNotExist:
            lens = Lens(library_id=library_id, name=lens_name, earliest_photo=date_taken,
                        latest_photo=date_taken)
            lens.save()

    photo = None
    if date_taken:
        try:
            # TODO: Match on file number/file name as well
            photo = Photo.objects.get(taken_at=date_taken)
        except Photo.DoesNotExist:
            pass

    latitude = None
    longitude = None
    if metadata.get('GPS Position'):
        latitude, longitude = parse_gps_location(metadata.get('GPS Position'))

    iso_speed = None
    if metadata.get('ISO'):
        try:
            iso_speed = int(re.search(r'[0-9]+', metadata.get('ISO')).group(0))
        except AttributeError:
            pass
    if not photo:
        # Save Photo
        aperture = None
        aperturestr = metadata.get('Aperture')
        if aperturestr:
            try:
                aperture = Decimal(aperturestr)
                if aperture.is_infinite():
                    aperture = None
            except:
                pass

        photo = Photo(
            library_id=library_id,
            taken_at=date_taken,
            taken_by=metadata.get('Artist') or None,
            aperture=aperture,
            exposure=metadata.get('Exposure Time') or None,
            iso_speed=iso_speed,
            focal_length=metadata.get('Focal Length') and metadata.get('Focal Length').split(' ', 1)[0] or None,
            flash=metadata.get('Flash') and 'on' in metadata.get('Flash').lower() or False,
            metering_mode=metadata.get('Metering Mode') or None,
            drive_mode=metadata.get('Drive Mode') or None,
            shooting_mode=metadata.get('Shooting Mode') or None,
            camera=camera,
            lens=lens,
            latitude=latitude,
            longitude=longitude,
            altitude=metadata.get('GPS Altitude') and metadata.get('GPS Altitude').split(' ')[0],
            star_rating=metadata.get('Rating')
        )
        photo.save()

        for subject in metadata.get('Subject', '').split(','):
            subject = subject.strip()
            if subject:
                tag = Tag.objects.create(library_id=library_id, name=subject, type="G")
                PhotoTag.objects.create(
                    photo=photo,
                    tag=tag,
                    confidence=1.0
            )
    else:
        for photo_file in photo.files.all():
            if not os.path.exists(photo_file.path):
                photo_file.delete()

    width = metadata.get('Image Width')
    height = metadata.get('Image Height')
    if metadata.get('Orientation') in ['Rotate 90 CW', 'Rotate 270 CCW', 'Rotate 90 CCW', 'Rotate 270 CW']:
        old_width = width
        width = height
        height = old_width

    # Save PhotoFile
    photo_file.photo = photo
    photo_file.path = path
    photo_file.width = width
    photo_file.height = height
    photo_file.mimetype = mimetypes.guess_type(path)[0]
    photo_file.file_modified_at = file_modified_at
    photo_file.bytes = os.stat(path).st_size
    photo_file.preferred = False  # TODO
    photo_file.save()

    # Create task to ensure JPEG version of file exists (used for thumbnailing, analysing etc.)
    Task(
        type='ensure_raw_processed',
        subject_id=photo.id,
        complete_with_children=True
    ).save()

    return photo

def delete_photo_record(photo_file_obj):
    """Delete photo record if photo not exixts on library path."""
    delete_photofile_and_photo_record(photo_file_obj)
    Tag.objects.filter(photo_tags=None).delete()
    Camera.objects.filter(photos=None).delete()
    Lens.objects.filter(photos=None).delete()
    return True

def move_or_rename_photo(photo_old_path, photo_new_path, library_id):
    """Rename a photoFile or change the path while moving photo in child directory."""
    try:
        photo_file = PhotoFile.objects.get(path=photo_old_path)
        photo_file.path = photo_new_path
        photo_file.save()
        return photo_file
    except Exception as e:
        return True

def delete_child_dir_all_photos(directory_path, library_id):
    """When a child directory deleted it delete all the photo records of that directory."""
    for photo_file_obj in PhotoFile.objects.filter(path__startswith=directory_path):
        delete_photofile_and_photo_record(photo_file_obj)
    Tag.objects.filter(photo_tags=None).delete()
    Camera.objects.filter(photos=None).delete()
    Lens.objects.filter(photos=None).delete()
    return True

def delete_photofile_and_photo_record(photo_file_obj):
    """Delete photoFile object with its photo object."""
    photo_obj = photo_file_obj.photo
    photo_file_obj.delete()
    if not photo_obj.files.all():
        photo_obj.delete()
<|endoftext|>


<commit_before>from errno import ENOENT as NoSuchFile
from keydb import KeyDB
from keydb import KeyDBWindow
from keydb import KeyDBFactory
from fs import Path
from fs import ensure_link, ensure_symlink, ensure_readonly
from snapshot import TreeSnapshot
from snapshot import KeySnapshot
from snapshot import snap_reduce
from os.path import sep
from itertools import combinations
from func_prototypes import typed, returned
from functools import partial
from farmfs.util import *
import re

def _metadata_path(root):
  assert isinstance(root, Path)
  return root.join(".farmfs")

@returned(Path)
@typed(Path)
def _keys_path(root):
  return _metadata_path(root).join("keys")

@returned(Path)
@typed(Path)
def _snaps_path(root):
  return _metadata_path(root).join("snaps")

def mkfs(root, udd):
  assert isinstance(root, Path)
  assert isinstance(udd, Path)
  root.mkdir()
  _metadata_path(root).mkdir()
  _keys_path(root).mkdir()
  _snaps_path(root).mkdir()
  kdb = KeyDB(_keys_path(root))
  # Make sure root key is removed.
  kdb.delete("root")
  kdb.write('udd', str(udd))
  udd.mkdir()
  kdb.write('status', {})
  vol = FarmFSVolume(root)

@returned(basestring)
@typed(basestring, int, int)
def _checksum_to_path(checksum, num_segs=3, seg_len=3):
  segs = [ checksum[i:i+seg_len] for i in range(0, min(len(checksum), seg_len * num_segs), seg_len)]
  segs.append(checksum[num_segs*seg_len:])
  return sep.join(segs)

_sep_replace_ = re.compile(sep)
@returned(basestring)
@typed(basestring)
def _remove_sep_(path):
    return _sep_replace_.subn("",path)[0]

def reverser(num_segs=3):
  r = re.compile("((\/([0-9]|[a-f])+){%d})$" % (num_segs+1))
  def checksum_from_link(link):
    m = r.search(str(link))
    if (m):
      csum_slash = m.group()[1:]
      csum = _remove_sep_(csum_slash)
      return csum
    else:
      raise ValueError("link %s checksum didn't parse" %(link))
  return checksum_from_link

def _validate_checksum(link2csum, path):
  csum = path.checksum()
  link_csum = link2csum(path)
  return csum == link_csum

def directory_signatures(snap):
  dirs = {}
  for entry in snap:
    if entry.is_link():
      (path, _, ref) = entry.get_tuple()
      parent = Path(path).parent()
      try:
        dirs[parent].update([ref])
      except KeyError:
        dirs[parent] = set([ref])
  return dirs

def encode_volume(vol):
  return str(vol.root)

def decode_volume(vol):
  return FarmFSVolume(Path(vol))

def encode_snapshot(snap):
  return map(lambda x: x.get_dict(), snap)

def decode_snapshot(splitter, reverser, data):
  return KeySnapshot(data, splitter, reverser)

class FarmFSVolume:
  def __init__(self, root):
    assert isinstance(root, Path)
    self.root = root
    self.mdd = _metadata_path(root)
    self.keydb = KeyDB(_keys_path(root))
    self.udd = Path(self.keydb.read('udd'))
    self.reverser = reverser()
    self.snapdb = KeyDBFactory(KeyDBWindow("snaps", self.keydb), encode_snapshot, partial(decode_snapshot, _checksum_to_path, self.reverser))
    self.remotedb = KeyDBFactory(KeyDBWindow("remotes", self.keydb), encode_volume, decode_volume)

    exclude_file = Path('.farmignore', self.root)
    self.exclude = [str(self.mdd)]
    try:
        with exclude_file.open('r') as exclude_fd:
          for pattern in exclude_fd.readlines():
            pattern = str(Path(pattern.strip(), root))
            self.exclude.append(pattern)
    except IOError as e:
      if e.errno == NoSuchFile:
          pass
      else: raise e

  def csum_to_name(self, csum):
    """Return string name of link relative to udd"""
    #TODO someday when csums are parameterized, we inject the has params here.
    return _checksum_to_path(csum)

  def csum_to_path(self, csum):
    """Return absolute Path to a blob given a csum"""
    return Path(self.csum_to_name(csum), self.udd)

  """Yield set of files not backed by FarmFS under path"""
  def thawed(self, path):
    for (entry, type_) in path.entries(self.exclude):
      if type_ == "file":
        yield entry

  """Yield set of files backed by FarmFS under path"""
  def frozen(self, path):
    for (entry, type_) in path.entries(self.exclude):
      if type_ == "link":
        yield entry

  #NOTE: This assumes a posix storage engine.
  def freeze(self, path):
    assert isinstance(path, Path)
    assert isinstance(self.udd, Path)
    csum = path.checksum()
    blob = self.csum_to_path(csum)
    assert blob == self.udd.join(_checksum_to_path(csum))
    duplicate = blob.exists()
    if not duplicate:
      ensure_link(blob, path)
      ensure_readonly(blob)
    ensure_symlink(path, blob)
    ensure_readonly(path)
    return {"path":path, "csum":csum, "was_dup":duplicate}

  #Note: This assumes a posix storage engine.
  def thaw(self, user_path):
    assert isinstance(user_path, Path)
    csum_path = user_path.readlink()
    user_path.unlink()
    csum_path.copy(user_path)
    return user_path

  """Find all broken links and point them back at UDD"""
  def repair_link(self, path):
    assert(path.islink())
    oldlink = path.readlink()
    if oldlink.isfile():
        print "Link %s is ok" % path
        return
    csum = self.reverser(oldlink)
    newlink = self.csum_to_path(csum)
    assert newlink == Path(_checksum_to_path(csum), self.udd)
    if not newlink.isfile():
      raise ValueError("%d is missing, cannot relink" % newlink)
    else:
      print "Relinking %s from %s to %s" % (path, oldlink, newlink)
      path.unlink()
      path.symlink(newlink)

  def check_userdata_hashes(self):#TODO MAKE A FUNCTOR
    """Make sure all backed file hashes match thier file contents"""
    link2csum = reverser()
    for (path, type_) in self.udd.entries():
      if type_ == "file":
        if not _validate_checksum(link2csum, path):
          yield path

  def check_link(self, udd_path):
    """Returns true if link is valid, false if invalid"""
    assert isinstance(udd_path, Path)
    return udd_path.exists();
  """
  TODO: There is a problem in checklinks where snaps and trees seem different.
  ('CORRUPTION: broken link in ', u'd41/d8c/d98/f00b204e9800998ecf8427e')
  ('CORRUPTION: broken link in ', u'/d41/d8c/d98/f00b204e9800998ecf8427e')
  """
  def check_links2(self):
    tree = self.tree()
    snaps = map(lambda x: self.snapdb.read(x), self.snapdb.list())
    select_links = partial(filter, lambda x: x.is_link())
    get_checksum = fmap(lambda x:x.csum())
    get_paths = fmap(self.csum_to_path)
    select_broken = partial(filter, lambda x: not x.exists())
    return transduce(
        concat,
        select_links,
        get_checksum,
        uniq,
        get_paths,
        select_broken,
        ) ([tree]+snaps) #TODO ADD SNAPS IN TOO.

  def check_links(self): #TODO MAKE A FUNCTOR.
    """Make sure that all links in the tree and in all snaps are backed in userdata"""
    for csum in self.count().keys(): #TODO we shouldn't get the list from count. its walking the tree wrong. If it hits a file, it crashes.
      path = self.csum_to_path(csum)
      if not self.check_link(path):
        yield path

  def fsck(self):
    for bad_link in self.check_links():
      yield "CORRUPTION: broken link in ", bad_link
    for bad_link in self.check_links2():
      yield "*** broken link in ", bad_link
    for bad_hash in self.check_userdata_hashes():
      yield "CORRUPTION: checksum mismatch in ", bad_hash

  """Get a snap object which represents the tree of the volume."""
  def tree(self):
    tree_snap = TreeSnapshot(self.root, self.udd, self.exclude, reverser=self.reverser)
    return tree_snap

  #TODO would be good to move this out of volume.
  #TODO would be good to turn this into a more composable design.
  def count(self):
    """Return a {checksum : count} for each unique file backed by FarmFS"""
    snaps = [self.tree()] + map(lambda x: self.snapdb.read(x), self.snapdb.list())
    counts = snap_reduce(snaps)
    return counts

  """Yields a set of paths which reference a given checksum_path name."""
  def reverse(self, udd_name):
    #TODO SCAN THE SNAPS FOR THIS SILLY PANTS.
    for (path, type_) in self.root.entries(self.exclude):
      if type_ == "link":
        ud_path = path.readlink()
        if ud_path == udd_name:
          yield path

  """ Yield all the relative paths (basestring) for all the files in the userdata store."""
  def userdata(self):
   # We populate counts with all hash paths from the userdata directory.
   for (path, type_) in self.udd.entries():
     assert isinstance(path, Path)
     if type_ == "file":
       yield path.relative_to(self.udd)
     elif type_ == "dir":
       pass
     else:
       raise ValueError("%s is f invalid type %s" % (path, type_))

  """Yields the names of files which are being garbage collected"""
  def gc(self):
    referenced_hashes = set(self.count().keys())
    udd_hashes = set(self.userdata())
    missing_data = referenced_hashes - udd_hashes
    assert len(missing_data) == 0, "Missing %s\nReferenced %s\nExisting %s\n" % (missing_data, referenced_hashes, udd_hashes)
    orphaned_data = udd_hashes - referenced_hashes
    for blob in orphaned_data:
      yield blob
      blob_path = self.udd.join(blob)
      blob_path.unlink(clean=self.udd)

  """Yields similarity data for directories"""
  def similarity(self):
    tree = self.tree()
    dir_sigs = directory_signatures(tree)
    combos = combinations(dir_sigs.items(), 2)
    for ((dir_a, sigs_a), (dir_b, sigs_b)) in combos:
      jac_sim = float(len(sigs_a.intersection(sigs_b)))/len(sigs_a.union(sigs_b))
      yield (dir_a, dir_b, jac_sim)
<commit_msg>Add todos. <commit_after>from errno import ENOENT as NoSuchFile
from keydb import KeyDB
from keydb import KeyDBWindow
from keydb import KeyDBFactory
from fs import Path
from fs import ensure_link, ensure_symlink, ensure_readonly
from snapshot import TreeSnapshot
from snapshot import KeySnapshot
from snapshot import snap_reduce
from os.path import sep
from itertools import combinations
from func_prototypes import typed, returned
from functools import partial
from farmfs.util import *
import re

def _metadata_path(root):
  assert isinstance(root, Path)
  return root.join(".farmfs")

@returned(Path)
@typed(Path)
def _keys_path(root):
  return _metadata_path(root).join("keys")

@returned(Path)
@typed(Path)
def _snaps_path(root):
  return _metadata_path(root).join("snaps")

def mkfs(root, udd):
  assert isinstance(root, Path)
  assert isinstance(udd, Path)
  root.mkdir()
  _metadata_path(root).mkdir()
  _keys_path(root).mkdir()
  _snaps_path(root).mkdir()
  kdb = KeyDB(_keys_path(root))
  # Make sure root key is removed.
  kdb.delete("root")
  kdb.write('udd', str(udd))
  udd.mkdir()
  kdb.write('status', {})
  vol = FarmFSVolume(root)

@returned(basestring)
@typed(basestring, int, int)
def _checksum_to_path(checksum, num_segs=3, seg_len=3):
  segs = [ checksum[i:i+seg_len] for i in range(0, min(len(checksum), seg_len * num_segs), seg_len)]
  segs.append(checksum[num_segs*seg_len:])
  return sep.join(segs)

_sep_replace_ = re.compile(sep)
@returned(basestring)
@typed(basestring)
def _remove_sep_(path):
    return _sep_replace_.subn("",path)[0]

def reverser(num_segs=3):
  r = re.compile("((\/([0-9]|[a-f])+){%d})$" % (num_segs+1))
  def checksum_from_link(link):
    m = r.search(str(link))
    if (m):
      csum_slash = m.group()[1:]
      csum = _remove_sep_(csum_slash)
      return csum
    else:
      raise ValueError("link %s checksum didn't parse" %(link))
  return checksum_from_link

def _validate_checksum(link2csum, path):
  csum = path.checksum()
  link_csum = link2csum(path)
  return csum == link_csum

def directory_signatures(snap):
  dirs = {}
  for entry in snap:
    if entry.is_link():
      (path, _, ref) = entry.get_tuple()
      parent = Path(path).parent()
      try:
        dirs[parent].update([ref])
      except KeyError:
        dirs[parent] = set([ref])
  return dirs

def encode_volume(vol):
  return str(vol.root)

def decode_volume(vol):
  return FarmFSVolume(Path(vol))

def encode_snapshot(snap):
  return map(lambda x: x.get_dict(), snap)

def decode_snapshot(splitter, reverser, data):
  return KeySnapshot(data, splitter, reverser)

class FarmFSVolume:
  def __init__(self, root):
    assert isinstance(root, Path)
    self.root = root
    self.mdd = _metadata_path(root)
    self.keydb = KeyDB(_keys_path(root))
    self.udd = Path(self.keydb.read('udd'))
    self.reverser = reverser()
    self.snapdb = KeyDBFactory(KeyDBWindow("snaps", self.keydb), encode_snapshot, partial(decode_snapshot, _checksum_to_path, self.reverser))
    self.remotedb = KeyDBFactory(KeyDBWindow("remotes", self.keydb), encode_volume, decode_volume)

    exclude_file = Path('.farmignore', self.root)
    self.exclude = [str(self.mdd)]
    try:
        with exclude_file.open('r') as exclude_fd:
          for pattern in exclude_fd.readlines():
            pattern = str(Path(pattern.strip(), root))
            self.exclude.append(pattern)
    except IOError as e:
      if e.errno == NoSuchFile:
          pass
      else: raise e

  def csum_to_name(self, csum):
    """Return string name of link relative to udd"""
    #TODO someday when csums are parameterized, we inject the has params here.
    return _checksum_to_path(csum)

  def csum_to_path(self, csum):
    """Return absolute Path to a blob given a csum"""
    return Path(self.csum_to_name(csum), self.udd)

  """Yield set of files not backed by FarmFS under path"""
  def thawed(self, path):
    for (entry, type_) in path.entries(self.exclude):
      if type_ == "file":
        yield entry

  """Yield set of files backed by FarmFS under path"""
  def frozen(self, path):
    for (entry, type_) in path.entries(self.exclude):
      if type_ == "link":
        yield entry

  #NOTE: This assumes a posix storage engine.
  def freeze(self, path):
    assert isinstance(path, Path)
    assert isinstance(self.udd, Path)
    csum = path.checksum()
    blob = self.csum_to_path(csum)
    assert blob == self.udd.join(_checksum_to_path(csum))
    duplicate = blob.exists()
    if not duplicate:
      ensure_link(blob, path)
      ensure_readonly(blob)
    ensure_symlink(path, blob)
    ensure_readonly(path)
    return {"path":path, "csum":csum, "was_dup":duplicate}

  #Note: This assumes a posix storage engine.
  def thaw(self, user_path):
    assert isinstance(user_path, Path)
    csum_path = user_path.readlink()
    user_path.unlink()
    csum_path.copy(user_path)
    return user_path

  """Find all broken links and point them back at UDD"""
  def repair_link(self, path):
    assert(path.islink())
    oldlink = path.readlink()
    if oldlink.isfile():
        print "Link %s is ok" % path
        return
    csum = self.reverser(oldlink)
    newlink = self.csum_to_path(csum)
    assert newlink == Path(_checksum_to_path(csum), self.udd)
    if not newlink.isfile():
      raise ValueError("%d is missing, cannot relink" % newlink)
    else:
      print "Relinking %s from %s to %s" % (path, oldlink, newlink)
      path.unlink()
      path.symlink(newlink)

  def check_userdata_hashes(self):#TODO MAKE A FUNCTOR
    """Make sure all backed file hashes match thier file contents"""
    link2csum = reverser()
    for (path, type_) in self.udd.entries():
      if type_ == "file":
        if not _validate_checksum(link2csum, path):
          yield path

  def check_link(self, udd_path):
    """Returns true if link is valid, false if invalid"""
    assert isinstance(udd_path, Path)
    return udd_path.exists();
  def check_links(self):
    #TODO message would be more helpful as {snap, path, csum}
    tree = self.tree()
    snaps = map(lambda x: self.snapdb.read(x), self.snapdb.list())
    select_links = partial(filter, lambda x: x.is_link())
    get_checksum = fmap(lambda x:x.csum())
    get_paths = fmap(self.csum_to_path)
    select_broken = partial(filter, lambda x: not x.exists())
    return transduce(
        concat,
        select_links,
        get_checksum,
        uniq,
        get_paths,
        select_broken,
        ) ([tree]+snaps) #TODO ADD SNAPS IN TOO.

  """
  ('CORRUPTION: broken link in ', u'd41/d8c/d98/f00b204e9800998ecf8427e')
  ('CORRUPTION: broken link in ', u'/d41/d8c/d98/f00b204e9800998ecf8427e')
  """
  def fsck(self):
    for bad_link in self.check_links():
      yield "CORRUPTION: broken link in ", bad_link
    for bad_hash in self.check_userdata_hashes():
      yield "CORRUPTION: checksum mismatch in ", bad_hash

  """Get a snap object which represents the tree of the volume."""
  def tree(self):
    tree_snap = TreeSnapshot(self.root, self.udd, self.exclude, reverser=self.reverser)
    return tree_snap

  #TODO would be good to move this out of volume.
  #TODO would be good to turn this into a more composable design.
  def count(self):
    """Return a {checksum : count} for each unique file backed by FarmFS"""
    snaps = [self.tree()] + map(lambda x: self.snapdb.read(x), self.snapdb.list())
    counts = snap_reduce(snaps)
    return counts

  """Yields a set of paths which reference a given checksum_path name."""
  def reverse(self, udd_name):
    #TODO SCAN THE SNAPS FOR THIS SILLY PANTS.
    for (path, type_) in self.root.entries(self.exclude):
      if type_ == "link":
        ud_path = path.readlink()
        if ud_path == udd_name:
          yield path

  """ Yield all the relative paths (basestring) for all the files in the userdata store."""
  def userdata(self):
   # We populate counts with all hash paths from the userdata directory.
   for (path, type_) in self.udd.entries():
     assert isinstance(path, Path)
     if type_ == "file":
       yield path.relative_to(self.udd)
     elif type_ == "dir":
       pass
     else:
       raise ValueError("%s is f invalid type %s" % (path, type_))

  """Yields the names of files which are being garbage collected"""
  def gc(self):
    referenced_hashes = set(self.count().keys())
    udd_hashes = set(self.userdata())
    missing_data = referenced_hashes - udd_hashes
    assert len(missing_data) == 0, "Missing %s\nReferenced %s\nExisting %s\n" % (missing_data, referenced_hashes, udd_hashes)
    orphaned_data = udd_hashes - referenced_hashes
    for blob in orphaned_data:
      yield blob
      blob_path = self.udd.join(blob)
      blob_path.unlink(clean=self.udd)

  """Yields similarity data for directories"""
  def similarity(self):
    tree = self.tree()
    dir_sigs = directory_signatures(tree)
    combos = combinations(dir_sigs.items(), 2)
    for ((dir_a, sigs_a), (dir_b, sigs_b)) in combos:
      jac_sim = float(len(sigs_a.intersection(sigs_b)))/len(sigs_a.union(sigs_b))
      yield (dir_a, dir_b, jac_sim)
<|endoftext|>


<commit_before>try:
    from cStringIO import StringIO      # Python 2
except ImportError:
    from io import StringIO
from python_terraform import *
import pytest
import os
import logging
import re
import shutil
from itertools import combinations

logging.basicConfig(level=logging.DEBUG)
root_logger = logging.getLogger()
current_path = os.path.dirname(os.path.realpath(__file__))

STRING_CASES = [
     [
         lambda x: x.generate_cmd_string('apply', 'the_folder',
                                         no_color=IsFlagged),
         "terraform apply -no-color the_folder"
     ],
     [
         lambda x: x.generate_cmd_string('push', 'path', vcs=True,
                                         token='token',
                                         atlas_address='url'),
         "terraform push -vcs=true -token=token -atlas-address=url path"
     ],
 ]

CMD_CASES = [
    ['method', 'expected_output', 'expected_ret_code', 'expected_logs'],
    [
        [
            lambda x: x.cmd('plan', 'var_to_output', no_color=IsFlagged, var={'test_var': 'test'}) ,
            "doesn't need to do anything",
            0,
            ''
        ],
        # try import aws instance
        [
            lambda x: x.cmd('import', 'aws_instance.foo', 'i-abcd1234', no_color=IsFlagged),
            '',
            1,
            'command: terraform import -no-color aws_instance.foo i-abcd1234'
        ]
    ]
]

@pytest.fixture(scope='function')
def fmt_test_file(request):
    target = os.path.join(current_path, 'bad_fmt', 'test.backup')
    orgin = os.path.join(current_path, 'bad_fmt', 'test.tf')
    shutil.copy(orgin,
                target)

    def td():
        shutil.move(target, orgin)

    request.addfinalizer(td)
    return

@pytest.fixture()
def string_logger(request):
    log_stream = StringIO()
    handler = logging.StreamHandler(log_stream)
    root_logger.addHandler(handler)

    def td():
        root_logger.removeHandler(handler)
        log_stream.close()

    request.addfinalizer(td)
    return lambda: str(log_stream.getvalue())

class TestTerraform(object):
    def teardown_method(self, method):
        """ teardown any state that was previously setup with a setup_method
        call.
        """

        def purge(dir, pattern):
            for f in os.listdir(dir):
                if re.search(pattern, f):
                    if os.path.isfile(f):
                        os.remove(os.path.join(dir, f))

        purge('.', '.tfstate')

    @pytest.mark.parametrize([
                 "method", "expected"
             ], STRING_CASES)
    def test_generate_cmd_string(self, method, expected):
        tf = Terraform(working_dir=current_path)
        result = method(tf)

        strs = expected.split()
        for s in strs:
            assert s in result

    @pytest.mark.parametrize(*CMD_CASES)
    def test_cmd(self, method, expected_output, expected_ret_code, expected_logs, string_logger):
        tf = Terraform(working_dir=current_path)
        ret, out, err = method(tf)
        logs = string_logger()
        logs = logs.replace('\n', '')
        assert expected_output in out
        assert expected_ret_code == ret
        assert expected_logs in logs

    @pytest.mark.parametrize(
        ("folder", "variables", "var_files", "expected_output", "options"),
        [
            ("var_to_output",
             {'test_var': 'test'}, None, "test_output=test", {}),
            ("var_to_output", {'test_list_var': ['c', 'd']}, None, "test_list_output=[c,d]", {}),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, None, "test_map_output={a=ab=bc=cd=d}", {}),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, 'var_to_output/test_map_var.json', "test_map_output={a=ab=bc=cd=de=ef=f}", {}),
            ("var_to_output", {}, None, "\x1b[0m\x1b[1m\x1b[32mApplycomplete!", {"no_color": IsNotFlagged})
        ])
    def test_apply(self, folder, variables, var_files, expected_output, options):
        tf = Terraform(working_dir=current_path, variables=variables, var_file=var_files)
        ret, out, err = tf.apply(folder, **options)
        assert ret == 0
        assert expected_output in out.replace('\n', '').replace(' ', '')
        assert err == ''

    @pytest.mark.parametrize(
        ['cmd', 'args', 'options'],
        [
            # bool value
            ('fmt', ['bad_fmt'], {'list': False, 'diff': False})
        ]
    )
    def test_options(self, cmd, args, options, fmt_test_file):
        tf = Terraform(working_dir=current_path)
        ret, out, err = getattr(tf, cmd)(*args, **options)
        assert ret == 0
        assert out == ''

    def test_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        tf.read_state_file()
        assert tf.tfstate.modules[0]['path'] == ['root']

    def test_pre_load_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        assert tf.tfstate.modules[0]['path'] == ['root']

    @pytest.mark.parametrize(
        ("folder", 'variables'),
        [
            ("var_to_output", {'test_var': 'test'})
        ]
    )
    def test_override_default(self, folder, variables):
        tf = Terraform(working_dir=current_path, variables=variables)
        ret, out, err = tf.apply(folder, var={'test_var': 'test2'},
                                 no_color=IsNotFlagged)
        out = out.replace('\n', '')
        assert '\x1b[0m\x1b[1m\x1b[32mApply' in out
        out = tf.output('test_output')
        assert 'test2' in out

    @pytest.mark.parametrize(
        ("param"),
        [
            ({}),
            ({'module': 'test2'}),
        ]
    )
    def test_output(self, param, string_logger):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        tf.apply('var_to_output')
        result = tf.output('test_output', **param)
        regex = re.compile('terraform output (-module=test2 -json|-json -module=test2) test_output')
        log_str = string_logger()
        if param:
            assert re.search(regex, log_str), log_str
        else:
            assert result == 'test'

    def test_destroy(self):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        ret, out, err = tf.destroy('var_to_output')
        assert ret == 0
        assert 'Destroy complete! Resources: 0 destroyed.' in out

    @pytest.mark.parametrize(
        ("plan", "variables", "expected_ret"),
        [
            ('vars_require_input', {}, 1)
        ]
    )
    def test_plan(self, plan, variables, expected_ret):
        tf = Terraform(working_dir=current_path, variables=variables)
        ret, out, err = tf.plan(plan)
        assert ret == expected_ret

    def test_fmt(self, fmt_test_file):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        ret, out, err = tf.fmt(diff=True)
        assert ret == 0

    def test_import(self, string_logger):
        tf = Terraform(working_dir=current_path)
        tf.import_cmd('aws_instance.foo', 'i-abc1234', no_color=IsFlagged)
        assert 'command: terraform import -no-color aws_instance.foo i-abc1234' in string_logger()
<commit_msg>clean up files<commit_after>try:
    from cStringIO import StringIO      # Python 2
except ImportError:
    from io import StringIO
from python_terraform import *
import pytest
import os
import logging
import re
import shutil
import fnmatch

logging.basicConfig(level=logging.DEBUG)
root_logger = logging.getLogger()
current_path = os.path.dirname(os.path.realpath(__file__))

STRING_CASES = [
     [
         lambda x: x.generate_cmd_string('apply', 'the_folder',
                                         no_color=IsFlagged),
         "terraform apply -no-color the_folder"
     ],
     [
         lambda x: x.generate_cmd_string('push', 'path', vcs=True,
                                         token='token',
                                         atlas_address='url'),
         "terraform push -vcs=true -token=token -atlas-address=url path"
     ],
 ]

CMD_CASES = [
    ['method', 'expected_output', 'expected_ret_code', 'expected_logs'],
    [
        [
            lambda x: x.cmd('plan', 'var_to_output', no_color=IsFlagged, var={'test_var': 'test'}) ,
            "doesn't need to do anything",
            0,
            ''
        ],
        # try import aws instance
        [
            lambda x: x.cmd('import', 'aws_instance.foo', 'i-abcd1234', no_color=IsFlagged),
            '',
            1,
            'command: terraform import -no-color aws_instance.foo i-abcd1234'
        ]
    ]
]

@pytest.fixture(scope='function')
def fmt_test_file(request):
    target = os.path.join(current_path, 'bad_fmt', 'test.backup')
    orgin = os.path.join(current_path, 'bad_fmt', 'test.tf')
    shutil.copy(orgin,
                target)

    def td():
        shutil.move(target, orgin)

    request.addfinalizer(td)
    return

@pytest.fixture()
def string_logger(request):
    log_stream = StringIO()
    handler = logging.StreamHandler(log_stream)
    root_logger.addHandler(handler)

    def td():
        root_logger.removeHandler(handler)
        log_stream.close()

    request.addfinalizer(td)
    return lambda: str(log_stream.getvalue())

class TestTerraform(object):
    def teardown_method(self, method):
        """ teardown any state that was previously setup with a setup_method
        call.
        """

        def purge(dir, pattern):
            # for f in os.listdir(dir):
            #     # if re.search(pattern, f):
            #     glob.glob()
            #         if os.path.isfile(f):
            #             os.remove(os.path.join(dir, f))
            for root, dirnames, filenames in os.walk(dir):
                for filename in fnmatch.filter(filenames, pattern):
                    f = os.path.join(root, filename)
                    os.remove(f)
                for dirname in fnmatch.filter(dirnames, pattern):
                    d = os.path.join(root, dirname)
                    shutil.rmtree(d)

        purge('.', '*.tfstate')
        purge('.', '*.terraform')

    @pytest.mark.parametrize([
                 "method", "expected"
             ], STRING_CASES)
    def test_generate_cmd_string(self, method, expected):
        tf = Terraform(working_dir=current_path)
        result = method(tf)

        strs = expected.split()
        for s in strs:
            assert s in result

    @pytest.mark.parametrize(*CMD_CASES)
    def test_cmd(self, method, expected_output, expected_ret_code, expected_logs, string_logger):
        tf = Terraform(working_dir=current_path)
        ret, out, err = method(tf)
        logs = string_logger()
        logs = logs.replace('\n', '')
        assert expected_output in out
        assert expected_ret_code == ret
        assert expected_logs in logs

    @pytest.mark.parametrize(
        ("folder", "variables", "var_files", "expected_output", "options"),
        [
            ("var_to_output",
             {'test_var': 'test'}, None, "test_output=test", {}),
            ("var_to_output", {'test_list_var': ['c', 'd']}, None, "test_list_output=[c,d]", {}),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, None, "test_map_output={a=ab=bc=cd=d}", {}),
            ("var_to_output", {'test_map_var': {"c": "c", "d": "d"}}, 'var_to_output/test_map_var.json', "test_map_output={a=ab=bc=cd=de=ef=f}", {}),
            ("var_to_output", {}, None, "\x1b[0m\x1b[1m\x1b[32mApplycomplete!", {"no_color": IsNotFlagged})
        ])
    def test_apply(self, folder, variables, var_files, expected_output, options):
        tf = Terraform(working_dir=current_path, variables=variables, var_file=var_files)
        ret, out, err = tf.apply(folder, **options)
        assert ret == 0
        assert expected_output in out.replace('\n', '').replace(' ', '')
        assert err == ''

    @pytest.mark.parametrize(
        ['cmd', 'args', 'options'],
        [
            # bool value
            ('fmt', ['bad_fmt'], {'list': False, 'diff': False})
        ]
    )
    def test_options(self, cmd, args, options, fmt_test_file):
        tf = Terraform(working_dir=current_path)
        ret, out, err = getattr(tf, cmd)(*args, **options)
        assert ret == 0
        assert out == ''

    def test_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        tf.read_state_file()
        assert tf.tfstate.modules[0]['path'] == ['root']

    def test_pre_load_state_data(self):
        cwd = os.path.join(current_path, 'test_tfstate_file')
        tf = Terraform(working_dir=cwd, state='tfstate.test')
        assert tf.tfstate.modules[0]['path'] == ['root']

    @pytest.mark.parametrize(
        ("folder", 'variables'),
        [
            ("var_to_output", {'test_var': 'test'})
        ]
    )
    def test_override_default(self, folder, variables):
        tf = Terraform(working_dir=current_path, variables=variables)
        ret, out, err = tf.apply(folder, var={'test_var': 'test2'},
                                 no_color=IsNotFlagged)
        out = out.replace('\n', '')
        assert '\x1b[0m\x1b[1m\x1b[32mApply' in out
        out = tf.output('test_output')
        assert 'test2' in out

    @pytest.mark.parametrize(
        ("param"),
        [
            ({}),
            ({'module': 'test2'}),
        ]
    )
    def test_output(self, param, string_logger):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        tf.apply('var_to_output')
        result = tf.output('test_output', **param)
        regex = re.compile('terraform output (-module=test2 -json|-json -module=test2) test_output')
        log_str = string_logger()
        if param:
            assert re.search(regex, log_str), log_str
        else:
            assert result == 'test'

    def test_destroy(self):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        ret, out, err = tf.destroy('var_to_output')
        assert ret == 0
        assert 'Destroy complete! Resources: 0 destroyed.' in out

    @pytest.mark.parametrize(
        ("plan", "variables", "expected_ret"),
        [
            ('vars_require_input', {}, 1)
        ]
    )
    def test_plan(self, plan, variables, expected_ret):
        tf = Terraform(working_dir=current_path, variables=variables)
        ret, out, err = tf.plan(plan)
        assert ret == expected_ret

    def test_fmt(self, fmt_test_file):
        tf = Terraform(working_dir=current_path, variables={'test_var': 'test'})
        ret, out, err = tf.fmt(diff=True)
        assert ret == 0

    def test_import(self, string_logger):
        tf = Terraform(working_dir=current_path)
        tf.import_cmd('aws_instance.foo', 'i-abc1234', no_color=IsFlagged)
        assert 'command: terraform import -no-color aws_instance.foo i-abc1234' in string_logger()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# flake8: noqa

"""Test CLI tool."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import os.path as op

from click.testing import CliRunner

from .._misc import _write_text

#------------------------------------------------------------------------------
# Test CLI tool
#------------------------------------------------------------------------------

def test_cli_empty(temp_user_dir):
    # NOTE: make the import after the temp_user_dir fixture, to avoid
    # loading any user plugin affecting the CLI.
    from ..cli import phy

    runner = CliRunner()

    result = runner.invoke(phy, [])
    assert result.exit_code == 0

    result = runner.invoke(phy, ['--version'])
    assert result.exit_code == 0
    assert result.output.startswith('phy,')

    result = runner.invoke(phy, ['--help'])
    assert result.exit_code == 0
    assert result.output.startswith('Usage: phy')

def test_cli_plugins(temp_user_dir):

    # Write a CLI plugin.
    cli_plugin = """
        import click
        from phy import IPlugin

        class MyPlugin(IPlugin):
            def attach_to_cli(self, cli):
                @cli.command()
                def hello():
                    click.echo("hello world")
    """
    path = op.join(temp_user_dir, 'plugins/hello.py')
    _write_text(path, cli_plugin)

    runner = CliRunner()

    # NOTE: make the import after the temp_user_dir fixture, to avoid
    # loading any user plugin affecting the CLI.
    from ..cli import phy

    result = runner.invoke(phy, ['--help'])
    assert result.exit_code == 0
    assert 'hello' in result.output
<commit_msg>Updated CLI tests<commit_after># -*- coding: utf-8 -*-
# flake8: noqa

"""Test CLI tool."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import os.path as op

from click.testing import CliRunner
from pytest import yield_fixture

from .._misc import _write_text

#------------------------------------------------------------------------------
# Test CLI tool
#------------------------------------------------------------------------------

@yield_fixture
def runner():
    yield CliRunner()

def test_cli_empty(temp_user_dir, runner):

    # NOTE: make the import after the temp_user_dir fixture, to avoid
    # loading any user plugin affecting the CLI.
    from ..cli import phy, load_cli_plugins
    load_cli_plugins(phy)

    result = runner.invoke(phy, [])
    assert result.exit_code == 0

    result = runner.invoke(phy, ['--version'])
    assert result.exit_code == 0
    assert result.output.startswith('phy,')

    result = runner.invoke(phy, ['--help'])
    assert result.exit_code == 0
    assert result.output.startswith('Usage: phy')

def test_cli_plugins(temp_user_dir, runner):

    # Write a CLI plugin.
    cli_plugin = """
        import click
        from phy import IPlugin

        class MyPlugin(IPlugin):
            def attach_to_cli(self, cli):
                @cli.command()
                def hello():
                    click.echo("hello world")
    """
    path = op.join(temp_user_dir, 'plugins/hello.py')
    _write_text(path, cli_plugin)

    # NOTE: make the import after the temp_user_dir fixture, to avoid
    # loading any user plugin affecting the CLI.
    from ..cli import phy, load_cli_plugins
    load_cli_plugins(phy)

    # The plugin should have added a new command.
    result = runner.invoke(phy, ['--help'])
    assert result.exit_code == 0
    assert 'hello' in result.output

    # The plugin should have added a new command.
    result = runner.invoke(phy, ['hello'])
    assert result.exit_code == 0
    assert result.output == 'hello world\n'
<|endoftext|>


<commit_before>import numpy as np

from c3nav.mapdata.render.engines.base3d import Base3DEngine

class STLEngine(Base3DEngine):
    facet_template = (b'  facet normal %f %f %f\n'
                      b'    outer loop\n'
                      b'      vertex %.3f %.3f %.3f\n'
                      b'      vertex %.3f %.3f %.3f\n'
                      b'      vertex %.3f %.3f %.3f\n'
                      b'    endloop\n'
                      b'  endfacet')

    def _create_facet(self, facet) -> bytes:
        return self.facet_template % tuple(facet.flatten())

    def render(self) -> bytes:
        facets = np.vstack(self.vertices)
        facets = np.hstack((np.cross(facets[:, 1]-facets[:, 0], facets[:, 2]-facets[:, 1]).reshape((-1, 1, 3))*1e11,
                            facets))
        return (b'solid c3nav_export\n' +
                b'\n'.join((self._create_facet(facet) for facet in facets)) +
                b'\nendsolid c3nav_export\n')
<commit_msg>our stl normals are long enough, drop the factor<commit_after>import numpy as np

from c3nav.mapdata.render.engines.base3d import Base3DEngine

class STLEngine(Base3DEngine):
    facet_template = (b'  facet normal %f %f %f\n'
                      b'    outer loop\n'
                      b'      vertex %.3f %.3f %.3f\n'
                      b'      vertex %.3f %.3f %.3f\n'
                      b'      vertex %.3f %.3f %.3f\n'
                      b'    endloop\n'
                      b'  endfacet')

    def _create_facet(self, facet) -> bytes:
        return self.facet_template % tuple(facet.flatten())

    def render(self) -> bytes:
        facets = np.vstack(self.vertices)
        facets = np.hstack((np.cross(facets[:, 1]-facets[:, 0], facets[:, 2]-facets[:, 1]).reshape((-1, 1, 3)),
                            facets))
        return (b'solid c3nav_export\n' +
                b'\n'.join((self._create_facet(facet) for facet in facets)) +
                b'\nendsolid c3nav_export\n')
<|endoftext|>


<commit_before>import numpy as np
import copy
from simtk import openmm, unit

default_functions = {
    'alchemical_sterics' : 'lambda',
    'alchemical_electrostatics' : 'lambda',
    'alchemical_bonds' : 'lambda',
    'alchemical_angles' : 'lambda',
    'alchemical_torsions' : 'lambda'
    }

class NCMCEngine(object):
    def __init__(self, temperature=300.0*unit.kelvin, functions=default_functions, nsteps=1, timestep=1.0*unit.femtoseconds, constraint_tolerance=None):
        """
        This is the base class for NCMC switching between two different systems.

        Arguments
        ---------
        temperature : simtk.unit.Quantity with units compatible with kelvin
            The temperature at which switching is to be run
        functions : dict of str:str, optional, default=default_functions
            functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
            controls how alchemical context parameter 'parameter' is switched
        nsteps : int, optional, default=1
            The number of steps to use for switching.
        timestep : simtk.unit.Quantity with units compatible with femtoseconds, optional, default=1*femtosecond
            The timestep to use for integration of switching velocity Verlet steps.
        constraint_tolerance : float, optional, default=None
            If not None, this relative constraint tolerance is used for position and velocity constraints.

        """
        self.temperature = temperature
        self.functions = copy.deepcopy(functions)
        self.nsteps = nsteps
        self.timestep = timestep
        self.constraint_tolerance = constraint_tolerance

    def _getAvailableParameters(self, system):
        """
        Return a list of available context parameters defined in the system

        Parameters
        ----------
        system : simtk.openmm.System
            The system for which available context parameters are to be determined

        Returns
        -------
        parameters : list of str
            The list of available context parameters in the system

        """
        parameters = list()
        for force_index in range(system.getNumForces()):
            force = system.getForce(force_index)
            if hasattr(force, 'getNumGlobalParameters'):
                for parameter_index in range(force.getNumGlobalParameters()):
                    parameters.append(force.getGlobalParameterName(parameter_index))
        return parameters

    def integrate(self, alchemical_system, initial_positions, direction='insert', platform=None):
        """
        Performs NCMC switching according to the provided

        Parameters
        ----------
        alchemical_system : simtk.openmm.System object
            alchemically-modified system with atoms to be eliminated
        initial_positions : [n, 3] numpy.ndarray
            positions of the atoms in the old system
        direction : str, optional, default='insert'
            Direction of alchemical switching:
                'insert' causes lambda to switch from 0 to 1 over nsteps steps of integration
                'delete' causes lambda to switch from 1 to 0 over nsteps steps of integration
        platform : simtk.openmm.Platform, optional, default=None
            If not None, this platform is used for integration.

        Returns
        -------
        final_positions : simtk.unit.Quantity of dimensions [nparticles,3] with units compatible with angstroms
            The final positions after `nsteps` steps of alchemical switching
        logP : float
            The log acceptance probability of the switch

        """
        # Make a list of available context parameters.

        # Select subset of switching functions based on which alchemical parameters are present in the system.
        available_parameters = self._getAvailableParameters(alchemical_system)
        functions = { parameter_name : self.functions[parameter_name] for parameter_name in self.functions if (parameter_name in available_parameters) }

        # Create an NCMC velocity Verlet integrator.
        integrator = NCMCAlchemicalIntegrator(self.temperature, alchemical_system, functions, nsteps=self.nsteps, timestep=self.timestep)
        # Set the constraint tolerance if specified.
        if self.constraint_tolerance is not None:
            integrator.setConstraintTolerance(self.constraint_tolerance)
        # Create a context on the specified platform.
        if platform is not None:
            context = openmm.Context(alchemical_system, integrator, platform)
        else:
            context = openmm.Context(alchemical_system, integrator)
        context.setPositions(initial_positions)
        # Set velocities to temperature and apply velocity constraints.
        context.setVelocitiesToTemperature(self.temperature)
        context.applyVelocityConstraints(integrator.getConstraintTolerance())
        # Only take a single integrator step since all switching steps are unrolled in NCMCAlchemicalIntegrator.
        integrator.step(1)
        # Store final positions and log acceptance probability.
        final_positions = context.getState(getPositions=True).getPositions(asNumpy=True)
        logP = integrator.getLogAcceptanceProbability()
        # Clean up.
        del context, integrator
        # Return
        return [final_positions, logP]

class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
    """
    Use NCMC switching to annihilate or introduce particles alchemically.

    TODO:
    ----
    * We may need to avoid unrolling integration steps.

    Examples
    --------

    Annihilate a Lennard-Jones particle

    >>> # Create an alchemically-perturbed test system
    >>> from openmmtools import testsystems
    >>> testsystem = testsystems.LennardJonesCluster()
    >>> from alchemy import AbsoluteAlchemicalFactory
    >>> alchemical_atoms = [0]
    >>> factory = AbsoluteAlchemicalFactory(testsystem.system, ligand_atoms=alchemical_atoms)
    >>> alchemical_system = factory.createPerturbedSystem()
    >>> # Create an NCMC switching integrator.
    >>> functions = { 'alchemical_sterics' : 't' }
    >>> ncmc_integrator = NCMCAlchemicalIntegrator(alchemical_system, functions, mode='delete')
    >>> # Create a Context
    >>> context = openmm.Context(alchemical_system, ncmc_integrator)
    >>> context.setPositions(testsystem.positions)
    >>> # Run the integrator
    >>> ncmc_integrator.step(1)
    >>> # Retrieve the log acceptance probability
    >>> log_ncmc = ncmc_integrator.log_ncmc

    Turn on an atom and its associated angles and torsions in alanine dipeptide

    >>> # Create an alchemically-perturbed test system
    >>> from openmmtools import testsystems
    >>> testsystem = testsystems.AlanineDipeptideVacuum()
    >>> from alchemy import AbsoluteAlchemicalFactory
    >>> alchemical_atoms = [0,1,2,3] # terminal methyl group
    >>> factory = AbsoluteAlchemicalFactory(testsystem.system, ligand_atoms=alchemical_atoms, alchemical_torsions=True, alchemical_angles=True, annihilate_sterics=True, annihilate_electrostatics=True)
    >>> alchemical_state = AlchemicalState(lambda_sterics=0, lambda_torsions=0, lambda_angles=0)
    >>> alchemical_system = factory.createPerturbedSystem(alchemical_state)
    >>> # Create an NCMC switching integrator.
    >>> functions = { 'lambda_sterics' : 't', 'lambda_electrostatics' : 't**0.5', 'lambda_torsions' : 't', 'lambda_angles' : 't**2' }
    >>> ncmc_integrator = NCMCAlchemicalIntegrator(alchemical_system, functions, mode='insert')
    >>> # Create a Context
    >>> context = openmm.Context(alchemical_system, ncmc_integrator)
    >>> context.setPositions(testsystem.positions)
    >>> # Run the integrator
    >>> ncmc_integrator.step(1)
    >>> # Retrieve the log acceptance probability
    >>> log_ncmc = ncmc_integrator.log_ncmc

    """

    def __init__(self, temperature, system, functions, nsteps=10, timestep=1.0*unit.femtoseconds, mode='insert'):
        """
        Initialize an NCMC switching integrator to annihilate or introduce particles alchemically.

        Parameters
        ----------
        temperature : simtk.unit.Quantity with units compatible with kelvin
            The temperature to use for computing the NCMC acceptance probability.
        system : simtk.openmm.System
            The system to be simulated.
        functions : dict of str : str
            functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
            controls how alchemical context parameter 'parameter' is switched
        nsteps : int, optional, default=10
            The number of switching timesteps per call to integrator.step(1).
        timestep : simtk.unit.Quantity with units compatible with femtoseconds
            The timestep to use for each NCMC step.
        mode : str, optional, default='insert'
            One of ['insert', 'delete'].
            For `insert`, the parameter 'lambda' is switched from 0 to 1.
            For `delete`, the parameter 'lambda' is switched from 1 to 0.

        Note that each call to integrator.step(1) executes the entire integration program; this should not be called with more than one step.

        A symmetric protocol is used, in which the protocol begins and ends with a velocity Verlet step.

        TODO:
        * Add a global variable that causes termination of future calls to step(1) after the first

        """
        if (nsteps < 1):
            raise Exception("'nsteps' must be >= 1")

        if mode not in ['insert', 'delete']:
            raise Exception("mode must be one of ['insert', 'delete']; was '%s' instead" % mode)

        super(NCMCAlchemicalIntegrator, self).__init__(timestep)

        self.addGlobalVariable('initial_total_energy', 0.0) # initial total energy (kinetic + potential)
        self.addGlobalVariable('final_total_energy', 0.0) # final total energy (kinetic + potential)
        self.addGlobalVariable('log_ncmc_acceptance_probability', 0.0) # log of NCMC acceptance probability

        if mode == 'insert':
            self.addGlobalVariable('lambda', 0.0) # parameter switched from 0 to 1 during course of integrating internal 'nsteps' of dynamics
        elif mode == 'delete':
            self.addGlobalVariable('lambda', 1.0) # parameter switched from 1 to 0 during course of integrating internal 'nsteps' of dynamics

        self.addPerDofVariable("x1", 0) # for velocity Verlet with constraints

        # Compute kT in natural openmm units.
        kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
        kT = kB * temperature
        kT = kT.value_in_unit_system(unit.md_unit_system)

        # Constrain initial positions and velocities.
        self.addConstrainPositions()
        self.addConstrainVelocities()

        # Store initial total energy.
        self.addComputeGlobal('initial_total_energy', 'ke + potential')

        #
        # Initial Velocity Verlet propagation step
        #

        self.addUpdateContextState()
        self.addComputePerDof("v", "v+0.5*dt*f/m")
        self.addComputePerDof("x", "x+dt*v")
        self.addComputePerDof("x1", "x")
        self.addConstrainPositions()
        self.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt")
        self.addConstrainVelocities()

        # Unroll loop over NCMC steps.
        for step in range(nsteps):
            #
            # Alchemical perturbation step
            #

            if nsteps==1:
                delta_lambda = 1.0
            else:
                delta_lambda = float(step)/float(nsteps-1)

            if mode == 'insert':
                self.addComputeGlobal('lambda', 'lambda + %f' % delta_lambda)
            elif mode == 'delete':
                self.addComputeGlobal('lambda', 'lambda - %f' % delta_lambda)

            # Update Context parameters according to provided functions.
            for context_parameter in functions:
                self.addComputeGlobal(context_parameter, functions[context_parameter])

            #
            # Velocity Verlet propagation step
            #

            self.addUpdateContextState()
            self.addComputePerDof("v", "v+0.5*dt*f/m")
            self.addComputePerDof("x", "x+dt*v")
            self.addComputePerDof("x1", "x")
            self.addConstrainPositions()
            self.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt")
            self.addConstrainVelocities()

        # Store final total energy.
        self.addComputeGlobal('final_total_energy', 'ke + potential')

        # Compute log acceptance probability.
        self.addComputeGlobal('log_ncmc_acceptance_probability', '(final_total_energy - initial_total_energy) / %f' % kT)

    def getLogAcceptanceProbability(self):
        return self.getGlobalVariableByName('log_ncmc_acceptance_probability')
<commit_msg>Fixed a bug in doctests for NCMC integrator.<commit_after>import numpy as np
import copy
from simtk import openmm, unit

default_functions = {
    'alchemical_sterics' : 'lambda',
    'alchemical_electrostatics' : 'lambda',
    'alchemical_bonds' : 'lambda',
    'alchemical_angles' : 'lambda',
    'alchemical_torsions' : 'lambda'
    }

class NCMCEngine(object):
    def __init__(self, temperature=300.0*unit.kelvin, functions=default_functions, nsteps=1, timestep=1.0*unit.femtoseconds, constraint_tolerance=None):
        """
        This is the base class for NCMC switching between two different systems.

        Arguments
        ---------
        temperature : simtk.unit.Quantity with units compatible with kelvin
            The temperature at which switching is to be run
        functions : dict of str:str, optional, default=default_functions
            functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
            controls how alchemical context parameter 'parameter' is switched
        nsteps : int, optional, default=1
            The number of steps to use for switching.
        timestep : simtk.unit.Quantity with units compatible with femtoseconds, optional, default=1*femtosecond
            The timestep to use for integration of switching velocity Verlet steps.
        constraint_tolerance : float, optional, default=None
            If not None, this relative constraint tolerance is used for position and velocity constraints.

        """
        self.temperature = temperature
        self.functions = copy.deepcopy(functions)
        self.nsteps = nsteps
        self.timestep = timestep
        self.constraint_tolerance = constraint_tolerance

    def _getAvailableParameters(self, system):
        """
        Return a list of available context parameters defined in the system

        Parameters
        ----------
        system : simtk.openmm.System
            The system for which available context parameters are to be determined

        Returns
        -------
        parameters : list of str
            The list of available context parameters in the system

        """
        parameters = list()
        for force_index in range(system.getNumForces()):
            force = system.getForce(force_index)
            if hasattr(force, 'getNumGlobalParameters'):
                for parameter_index in range(force.getNumGlobalParameters()):
                    parameters.append(force.getGlobalParameterName(parameter_index))
        return parameters

    def integrate(self, alchemical_system, initial_positions, direction='insert', platform=None):
        """
        Performs NCMC switching according to the provided

        Parameters
        ----------
        alchemical_system : simtk.openmm.System object
            alchemically-modified system with atoms to be eliminated
        initial_positions : [n, 3] numpy.ndarray
            positions of the atoms in the old system
        direction : str, optional, default='insert'
            Direction of alchemical switching:
                'insert' causes lambda to switch from 0 to 1 over nsteps steps of integration
                'delete' causes lambda to switch from 1 to 0 over nsteps steps of integration
        platform : simtk.openmm.Platform, optional, default=None
            If not None, this platform is used for integration.

        Returns
        -------
        final_positions : simtk.unit.Quantity of dimensions [nparticles,3] with units compatible with angstroms
            The final positions after `nsteps` steps of alchemical switching
        logP : float
            The log acceptance probability of the switch

        """
        # Make a list of available context parameters.

        # Select subset of switching functions based on which alchemical parameters are present in the system.
        available_parameters = self._getAvailableParameters(alchemical_system)
        functions = { parameter_name : self.functions[parameter_name] for parameter_name in self.functions if (parameter_name in available_parameters) }

        # Create an NCMC velocity Verlet integrator.
        integrator = NCMCAlchemicalIntegrator(self.temperature, alchemical_system, functions, nsteps=self.nsteps, timestep=self.timestep)
        # Set the constraint tolerance if specified.
        if self.constraint_tolerance is not None:
            integrator.setConstraintTolerance(self.constraint_tolerance)
        # Create a context on the specified platform.
        if platform is not None:
            context = openmm.Context(alchemical_system, integrator, platform)
        else:
            context = openmm.Context(alchemical_system, integrator)
        context.setPositions(initial_positions)
        # Set velocities to temperature and apply velocity constraints.
        context.setVelocitiesToTemperature(self.temperature)
        context.applyVelocityConstraints(integrator.getConstraintTolerance())
        # Only take a single integrator step since all switching steps are unrolled in NCMCAlchemicalIntegrator.
        integrator.step(1)
        # Store final positions and log acceptance probability.
        final_positions = context.getState(getPositions=True).getPositions(asNumpy=True)
        logP = integrator.getLogAcceptanceProbability()
        # Clean up.
        del context, integrator
        # Return
        return [final_positions, logP]

class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
    """
    Use NCMC switching to annihilate or introduce particles alchemically.

    TODO:
    ----
    * We may need to avoid unrolling integration steps.

    Examples
    --------

    Annihilate a Lennard-Jones particle

    >>> # Create an alchemically-perturbed test system
    >>> from openmmtools import testsystems
    >>> testsystem = testsystems.LennardJonesCluster()
    >>> from alchemy import AbsoluteAlchemicalFactory
    >>> alchemical_atoms = [0]
    >>> factory = AbsoluteAlchemicalFactory(testsystem.system, ligand_atoms=alchemical_atoms)
    >>> alchemical_system = factory.createPerturbedSystem()
    >>> # Create an NCMC switching integrator.
    >>> functions = { 'alchemical_sterics' : 'lambda' }
    >>> ncmc_integrator = NCMCAlchemicalIntegrator(alchemical_system, functions, mode='delete')
    >>> # Create a Context
    >>> context = openmm.Context(alchemical_system, ncmc_integrator)
    >>> context.setPositions(testsystem.positions)
    >>> # Run the integrator
    >>> ncmc_integrator.step(1)
    >>> # Retrieve the log acceptance probability
    >>> log_ncmc = ncmc_integrator.log_ncmc

    Turn on an atom and its associated angles and torsions in alanine dipeptide

    >>> # Create an alchemically-perturbed test system
    >>> from openmmtools import testsystems
    >>> testsystem = testsystems.AlanineDipeptideVacuum()
    >>> from alchemy import AbsoluteAlchemicalFactory
    >>> alchemical_atoms = [0,1,2,3] # terminal methyl group
    >>> factory = AbsoluteAlchemicalFactory(testsystem.system, ligand_atoms=alchemical_atoms, alchemical_torsions=True, alchemical_angles=True, annihilate_sterics=True, annihilate_electrostatics=True)
    >>> alchemical_state = AlchemicalState(lambda_sterics=0, lambda_torsions=0, lambda_angles=0)
    >>> alchemical_system = factory.createPerturbedSystem(alchemical_state)
    >>> # Create an NCMC switching integrator.
    >>> functions = { 'lambda_sterics' : 'lambda', 'lambda_electrostatics' : 'lambda**0.5', 'lambda_torsions' : 'lambda', 'lambda_angles' : 'lambda**2' }
    >>> ncmc_integrator = NCMCAlchemicalIntegrator(alchemical_system, functions, mode='insert')
    >>> # Create a Context
    >>> context = openmm.Context(alchemical_system, ncmc_integrator)
    >>> context.setPositions(testsystem.positions)
    >>> # Run the integrator
    >>> ncmc_integrator.step(1)
    >>> # Retrieve the log acceptance probability
    >>> log_ncmc = ncmc_integrator.log_ncmc

    """

    def __init__(self, temperature, system, functions, nsteps=10, timestep=1.0*unit.femtoseconds, mode='insert'):
        """
        Initialize an NCMC switching integrator to annihilate or introduce particles alchemically.

        Parameters
        ----------
        temperature : simtk.unit.Quantity with units compatible with kelvin
            The temperature to use for computing the NCMC acceptance probability.
        system : simtk.openmm.System
            The system to be simulated.
        functions : dict of str : str
            functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
            controls how alchemical context parameter 'parameter' is switched
        nsteps : int, optional, default=10
            The number of switching timesteps per call to integrator.step(1).
        timestep : simtk.unit.Quantity with units compatible with femtoseconds
            The timestep to use for each NCMC step.
        mode : str, optional, default='insert'
            One of ['insert', 'delete'].
            For `insert`, the parameter 'lambda' is switched from 0 to 1.
            For `delete`, the parameter 'lambda' is switched from 1 to 0.

        Note that each call to integrator.step(1) executes the entire integration program; this should not be called with more than one step.

        A symmetric protocol is used, in which the protocol begins and ends with a velocity Verlet step.

        TODO:
        * Add a global variable that causes termination of future calls to step(1) after the first

        """
        if (nsteps < 1):
            raise Exception("'nsteps' must be >= 1")

        if mode not in ['insert', 'delete']:
            raise Exception("mode must be one of ['insert', 'delete']; was '%s' instead" % mode)

        super(NCMCAlchemicalIntegrator, self).__init__(timestep)

        self.addGlobalVariable('initial_total_energy', 0.0) # initial total energy (kinetic + potential)
        self.addGlobalVariable('final_total_energy', 0.0) # final total energy (kinetic + potential)
        self.addGlobalVariable('log_ncmc_acceptance_probability', 0.0) # log of NCMC acceptance probability

        if mode == 'insert':
            self.addGlobalVariable('lambda', 0.0) # parameter switched from 0 to 1 during course of integrating internal 'nsteps' of dynamics
        elif mode == 'delete':
            self.addGlobalVariable('lambda', 1.0) # parameter switched from 1 to 0 during course of integrating internal 'nsteps' of dynamics

        self.addPerDofVariable("x1", 0) # for velocity Verlet with constraints

        # Compute kT in natural openmm units.
        kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
        kT = kB * temperature
        kT = kT.value_in_unit_system(unit.md_unit_system)

        # Constrain initial positions and velocities.
        self.addConstrainPositions()
        self.addConstrainVelocities()

        # Store initial total energy.
        self.addComputeGlobal('initial_total_energy', 'ke + potential')

        #
        # Initial Velocity Verlet propagation step
        #

        self.addUpdateContextState()
        self.addComputePerDof("v", "v+0.5*dt*f/m")
        self.addComputePerDof("x", "x+dt*v")
        self.addComputePerDof("x1", "x")
        self.addConstrainPositions()
        self.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt")
        self.addConstrainVelocities()

        # Unroll loop over NCMC steps.
        for step in range(nsteps):
            #
            # Alchemical perturbation step
            #

            if nsteps==1:
                delta_lambda = 1.0
            else:
                delta_lambda = float(step)/float(nsteps-1)

            if mode == 'insert':
                self.addComputeGlobal('lambda', 'lambda + %f' % delta_lambda)
            elif mode == 'delete':
                self.addComputeGlobal('lambda', 'lambda - %f' % delta_lambda)

            # Update Context parameters according to provided functions.
            for context_parameter in functions:
                self.addComputeGlobal(context_parameter, functions[context_parameter])

            #
            # Velocity Verlet propagation step
            #

            self.addUpdateContextState()
            self.addComputePerDof("v", "v+0.5*dt*f/m")
            self.addComputePerDof("x", "x+dt*v")
            self.addComputePerDof("x1", "x")
            self.addConstrainPositions()
            self.addComputePerDof("v", "v+0.5*dt*f/m+(x-x1)/dt")
            self.addConstrainVelocities()

        # Store final total energy.
        self.addComputeGlobal('final_total_energy', 'ke + potential')

        # Compute log acceptance probability.
        self.addComputeGlobal('log_ncmc_acceptance_probability', '(final_total_energy - initial_total_energy) / %f' % kT)

    def getLogAcceptanceProbability(self):
        return self.getGlobalVariableByName('log_ncmc_acceptance_probability')
<|endoftext|>