<commit_before>import json
import os
import sys


FRAGMENTDIR = "fragments/"


class ValidationError(Exception):
  pass

class LoadError(Exception):
  pass


class Fragment:
  def __init__(self, atb_id):
    self.atb_id = atb_id
    self.atoms = []

  def add_atom(self, id, charge, other_id):
    self.atoms.append(Atom(id, charge, other_id))

  @property
  def __dict__(self):
    return {
      'atb_id': self.atb_id,
      'atoms': map(lambda a: a.__dict__, self.atoms)
    }


class Atom:
  def __init__(self, id, charge, other_id):
    self.id = id
    self.charge = charge
    self.other_id = other_id

  @property
  def __dict__(self):
    return {
      'id': self.id,
      'charge': self.charge,
      'other_id': self.other_id
    }


def load_off(off_name):
  if not os.path.exists(off_name):
    raise LoadError("Could not find fragment file %s" % off_name)

  with open(off_name, 'r') as fp:
    data = fp.read()

  try:
    return json.loads(data)
  except ValueError as e:
    raise LoadError("Invalid fragment file: %s" % e)

def get_fragments(off_name, needle):
  off = load_off(off_name)
  fragments = []
  for molecule in off["molecules"]:
    for fragment in molecule["fragments"]:
      if not "pairs" in fragment:
        continue

      aids = map(lambda p: p["id1"], fragment["pairs"])
      match = True
      for aid in needle:
        if not aid in aids:
          match = False
          break
      if match:
        frag = Fragment(molecule["atb_id"])
        for pair in fragment["pairs"]:
          frag.add_atom(pair["id1"], pair["charge"], pair["id2"])
        fragments.append(frag)
  return fragments

def find_fragments(args):
  try:
    data = json.loads(args)
  except ValueError as e:
    return {'error': "Query not in JSON format: %s" % e}

  try:
    validate_query(data)
  except ValidationError as e:
    return {'error': "Invalid query: %s" % e}

  # This is safe now, as all has been validated
  off_name = data["off"]
  needle = data["needle"]

  try:
    fragments = get_fragments(off_name, needle)
  except LoadError as e:
    return {'error': "Could not load fragments: %s" % e}

  return {'fragments': fragments}

def validate_query(data):
  if not 'off' in data or len(data['off']) == 0:
    raise ValidationError("FFID not set")
  elif not 'needle' in data or len(data['needle']) == 0:
    raise ValidationError("Needle not set")


if __name__ == "__main__":
  result = find_fragments(sys.argv[1])
  print json.dumps(result, default=lambda o: o.__dict__)<commit_msg>Fixed error message for missing OFF.<commit_after>import json
import os
import sys


FRAGMENTDIR = "fragments/"


class ValidationError(Exception):
  pass

class LoadError(Exception):
  pass


class Fragment:
  def __init__(self, atb_id):
    self.atb_id = atb_id
    self.atoms = []

  def add_atom(self, id, charge, other_id):
    self.atoms.append(Atom(id, charge, other_id))

  @property
  def __dict__(self):
    return {
      'atb_id': self.atb_id,
      'atoms': map(lambda a: a.__dict__, self.atoms)
    }


class Atom:
  def __init__(self, id, charge, other_id):
    self.id = id
    self.charge = charge
    self.other_id = other_id

  @property
  def __dict__(self):
    return {
      'id': self.id,
      'charge': self.charge,
      'other_id': self.other_id
    }


def load_off(off_name):
  if not os.path.exists(off_name):
    raise LoadError("Could not find fragment file %s" % off_name)

  with open(off_name, 'r') as fp:
    data = fp.read()

  try:
    return json.loads(data)
  except ValueError as e:
    raise LoadError("Invalid fragment file: %s" % e)

def get_fragments(off_name, needle):
  off = load_off(off_name)
  fragments = []
  for molecule in off["molecules"]:
    for fragment in molecule["fragments"]:
      if not "pairs" in fragment:
        continue

      aids = map(lambda p: p["id1"], fragment["pairs"])
      match = True
      for aid in needle:
        if not aid in aids:
          match = False
          break
      if match:
        frag = Fragment(molecule["atb_id"])
        for pair in fragment["pairs"]:
          frag.add_atom(pair["id1"], pair["charge"], pair["id2"])
        fragments.append(frag)
  return fragments

def find_fragments(args):
  try:
    data = json.loads(args)
  except ValueError as e:
    return {'error': "Query not in JSON format: %s" % e}

  try:
    validate_query(data)
  except ValidationError as e:
    return {'error': "Invalid query: %s" % e}

  # This is safe now, as all has been validated
  off_name = data["off"]
  needle = data["needle"]

  try:
    fragments = get_fragments(off_name, needle)
  except LoadError as e:
    return {'error': "Could not load fragments: %s" % e}

  return {'fragments': fragments}

def validate_query(data):
  if not 'off' in data or len(data['off']) == 0:
    raise ValidationError("OFF not set")
  elif not 'needle' in data or len(data['needle']) == 0:
    raise ValidationError("Needle not set")


if __name__ == "__main__":
  result = find_fragments(sys.argv[1])
  print json.dumps(result, default=lambda o: o.__dict__)
{'fraction_non_alphanumeric': 0.1253454401894986, 'fraction_numerical': 0.002368732727990525, 'mean_word_length': 3.102834008097166, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5899706', 'n_tokens_mistral': 1755, 'n_tokens_neox': 1710, 'n_words': 500}