<commit_before>#!/usr/bin/env python

import argparse
import curses
import curses.ascii
import os
import os.path
import string
import sys
import tempfile
import text_unidecode
import typing
import urllib.parse
import urllib.request
import validators

# constants

__version__ = "0.0.3"

DESCRIPTION = "{} is a tool to learn about something just copying it " \
              "by hand. Type Esc or ^[ to exit while running it." \
              .format(os.path.basename(__file__))

TTY_DEVICE_FILE = "/dev/tty" # POSIX compliant
UTF8 = "UTF-8"

QUIT_CHARS = {chr(curses.ascii.ESC), curses.ascii.ctrl('[')}
DELETE_CHARS = {chr(curses.ascii.DEL), chr(curses.ascii.BS),
                chr(curses.KEY_BACKSPACE), chr(curses.KEY_DC)}
CLEAR_CHAR = curses.ascii.ctrl('u')
CHEAT_CHAR = curses.ascii.ctrl('n')

CAN_CHEAT = False
SPACES_PER_TAB = 2

ATTR_CORRECT = curses.A_NORMAL
ATTR_ERROR = curses.A_REVERSE

SUPPORTED_SCHEMES = {"http", "https"}

# the order of bugs

def message(*text):
  print("{}:".format(os.path.basename(__file__)), *text, file=sys.stderr)

def error(*text):
  message("error:", *text)

def fail(*text):
  error(*text)
  exit(1)

# classes

class ConsoleApi:
  def __init__(self, window):
    window.keypad(True)
    window.scrollok(True)
    self.__window = window

  def __save_position(method):
    def wrapper(self, *args, **keyword_args):
      position = self.__window.getyx()
      method(self, *args, **keyword_args)
      self.__window.move(*position)
    return wrapper

  @property
  def screen_height(self):
    return self.__window.getmaxyx()[0]

  @property
  def screen_width(self):
    return self.__window.getmaxyx()[1]

  def get_char(self) -> str:
    return chr(self.__window.getch())

  @__save_position
  def put_char(self, char: str, attr=curses.A_NORMAL):
    assert self.__is_single_width_ascii_char(char)
    self.__window.addch(ord(char), attr)

  @__save_position
  def put_line(self, line_pos, text):
    assert len(text) <= self.screen_width
    self.__window.move(line_pos, 0)
    self.__window.clrtoeol()
    self.__window.addstr(line_pos, 0, text)

  def move(self, y, x):
    self.__window.move(y, x)

  def move_right(self):
    y, x = self.__window.getyx()
    if x == self.screen_width - 1: return
    self.__window.move(y, x + 1)

  def move_left(self):
    y, x = self.__window.getyx()
    if x == 0: return
    self.__window.move(y, x - 1)

  @__save_position
  def clear(self):
    self.__window.clear()

  @__save_position
  def scroll(self):
    self.__window.scroll()

  @staticmethod
  def __is_single_width_ascii_char(char: str):
    return len(char) == 1 and (curses.ascii.isprint(char) or char == " ")

class TypingGame:
  def __init__(self, api, example_file):
    self.__api = api
    self.__example_text = FormattedText(example_file,
                                        line_length=(api.screen_width - 1))
    self.__geometry = Geometry(api.screen_height)
    self.__input_text = ""
    if self.__example_text[0] == None:
      raise Exception("No line can be read from example source.")

  def play(self):
    self.__initialize_screen()

    game_over = False
    while not game_over:
      char = self.__api.get_char()

      if char in QUIT_CHARS:
        return
      elif char == CLEAR_CHAR:
        self.__clear_input_text()
      elif char in DELETE_CHARS:
        self.__delete_char()
      elif char == CHEAT_CHAR and CAN_CHEAT:
        game_over = self.__cheat()
      elif len(char) == 1 and (curses.ascii.isprint(char)
           or curses.ascii.isspace(char)):
        game_over = self.__add_char(char)

  def __add_char(self, char: str) -> bool:
    go_to_next_line = char == '\n' \
                      and self.__input_text == self.__example_text[0]
    if go_to_next_line and self.__example_text[1] == None:
      return True

    if go_to_next_line:
      self.__scroll()
    elif len(self.__input_text) == len(self.__example_text[0]):
      pass
    elif char == '\t':
      for _ in range(SPACES_PER_TAB):
        self.__add_char(' ')
    elif curses.ascii.isprint(char):
      self.__input_text += char
      attr = ATTR_CORRECT \
             if char == self.__example_text[0][len(self.__input_text) - 1] \
             else ATTR_ERROR
      self.__api.put_char(char, attr=attr)
      self.__api.move_right()
    return False

  def __delete_char(self):
    if len(self.__input_text) == 0: return
    self.__input_text = self.__input_text[:-1]
    self.__api.move_left()
    self.__api.put_char(self.__example_text[0][len(self.__input_text)])

  def __cheat(self) -> bool:
    if self.__example_text[1] == None:
      return True
    self.__scroll()
    return False

  def __initialize_screen(self):
    self.__api.clear()
    self.__print_all_example_text()
    self.__clear_input_text()

  def __scroll(self):
    self.__api.scroll()
    del self.__example_text[0]
    self.__print_bottom_example_text()
    self.__clear_input_text()

  def __clear_input_text(self):
    self.__input_text = ""
    self.__api.put_line(self.__geometry.current_line, self.__example_text[0])
    self.__api.move(self.__geometry.current_line, 0)

  def __print_bottom_example_text(self):
    index = self.__geometry.bottom_line - self.__geometry.current_line
    if self.__example_text[index] != None:
      self.__api.put_line(self.__geometry.bottom_line,
                          self.__example_text[index])

  def __print_all_example_text(self):
    for index in range(self.__geometry.bottom_line
                       - self.__geometry.current_line + 1):
      if self.__example_text[index] == None: break
      self.__api.put_line(self.__geometry.current_line + index,
                          self.__example_text[index])

class Geometry:
  def __init__(self, screen_height):
    self.current_line = (screen_height - 1) // 2
    self.bottom_line = screen_height - 1

class FormattedText:
  def __init__(self, text_file, line_length=79):
    self.__file = text_file
    self.__line_length = line_length
    self.__buffered_lines = []

  def __del__(self):
    self.__file.close()

  def __getitem__(self, index: int) -> str:
    self.__read_lines_from_file()

    return self.__buffered_lines[index] if index < len(self.__buffered_lines) \
           else None

  def __delitem__(self, index: int):
    del self.__buffered_lines[index]

  def __read_lines_from_file(self):
    lines = self.__file.readlines()

    for line in lines:
      self.__buffer_line(self.__preprocess_text(line))

  def __buffer_line(self, line):
    line = line.rstrip()

    while len(line) > self.__line_length:
      buffered_line, line = self.__split_line(line)
      self.__buffered_lines.append(buffered_line)
    self.__buffered_lines.append(line)

  def __split_line(self, line):
    return line[:self.__line_length].rstrip(), \
           line[self.__line_length:].rstrip()

  @staticmethod
  def __preprocess_text(text):
    return "".join(char for char in text_unidecode.unidecode(
                   text.replace('\t', ' ' * SPACES_PER_TAB))
                   if char in string.printable)

# functions

def reset_stdin():
  TEMPORARY_FD = 3

  os.dup2(sys.stdin.fileno(), TEMPORARY_FD)
  os.close(sys.stdin.fileno())
  sys.stdin = open(TTY_DEVICE_FILE)

  return os.fdopen(TEMPORARY_FD)

def initialize_curses():
  window = curses.initscr()
  curses.noecho()
  curses.cbreak()
  curses.start_color()
  curses.use_default_colors()
  return window

def finalize_curses():
  curses.nocbreak()
  curses.echo()
  curses.endwin() # should be at the very last line

def parse_args():
  arg_parser = argparse.ArgumentParser(description=DESCRIPTION)
  arg_parser.add_argument("example_path", nargs='?', default=None,
                          help="path or URI to example")
  arg_parser.add_argument("-c", "--cheat",
                          dest="can_cheat",
                          action="store_true",
                          help="enable the cheat key, {}"
                               .format(curses.ascii.unctrl(CHEAT_CHAR)))
  arg_parser.add_argument("-t", "--spaces-per-tab",
                          type=int, dest="spaces_per_tab",
                          help="set number of spaces per tab")
  arg_parser.add_argument("-v", "--version",
                          dest="show_version",
                          action="store_true", default=False,
                          help="show version information")

  args = arg_parser.parse_args()

  if args.show_version:
    print("version:", __version__)
    exit()

  if args.spaces_per_tab != None:
    global SPACES_PER_TAB
    SPACES_PER_TAB = args.spaces_per_tab
  if args.can_cheat == True:
    global CAN_CHEAT
    CAN_CHEAT = True

  return args

def get_example_file():
  example_path = parse_args().example_path

  if example_path != None and validators.url(example_path) \
      and sys.stdin.isatty():
    return get_remote_file(example_path)
  elif example_path != None and sys.stdin.isatty():
    return open(example_path)
  elif example_path == None and not sys.stdin.isatty():
    return reset_stdin()
  else:
    fail("Example is read from either the path or URI specified "
         "in the argument or stdin.")

def get_remote_file(uri):
  if urllib.parse.urlparse(uri).scheme not in SUPPORTED_SCHEMES:
    fail("Invalid scheme of URI is detected. "
         "(supported schemes: {})".format(", ".join(SUPPORTED_SCHEMES)))

  message("Loading page...")
  temporary_file = tempfile.TemporaryFile(mode="w+")
  with urllib.request.urlopen(uri) as response:
    temporary_file.write(response.read().decode(UTF8, "replace"))
  temporary_file.seek(0)
  return temporary_file

# main routine

def main():
  if not sys.stdout.isatty(): fail("stdout is not a tty.")

  example_file = get_example_file()

  try:
    # CAUTION:
    # You need to raise some Exception() instead of calling exit() here
    # to prevent curses from messing up your terminal.

    window = initialize_curses()

    TypingGame(ConsoleApi(window), example_file).play()
  finally:
    finalize_curses()

if __name__ == "__main__":
  main()
<commit_msg>Delete tab characters properly<commit_after>#!/usr/bin/env python

import argparse
import curses
import curses.ascii
import os
import os.path
import string
import sys
import tempfile
import text_unidecode
import typing
import unicodedata
import urllib.parse
import urllib.request
import validators

# constants

__version__ = "0.0.3"

DESCRIPTION = "{} is a tool to learn about something just copying it " \
              "by hand. Type Esc or ^[ to exit while running it." \
              .format(os.path.basename(__file__))

TTY_DEVICE_FILE = "/dev/tty" # POSIX compliant
UTF8 = "UTF-8"

QUIT_CHARS = {chr(curses.ascii.ESC), curses.ascii.ctrl('[')}
DELETE_CHARS = {chr(curses.ascii.DEL), chr(curses.ascii.BS),
                chr(curses.KEY_BACKSPACE), chr(curses.KEY_DC)}
CLEAR_CHAR = curses.ascii.ctrl('u')
CHEAT_CHAR = curses.ascii.ctrl('n')

CAN_CHEAT = False
SPACES_PER_TAB = 2

ATTR_CORRECT = curses.A_NORMAL
ATTR_ERROR = curses.A_REVERSE

SUPPORTED_SCHEMES = {"http", "https"}

# the order of bugs

def message(*text):
  print("{}:".format(os.path.basename(__file__)), *text, file=sys.stderr)

def error(*text):
  message("error:", *text)

def fail(*text):
  error(*text)
  exit(1)

# classes

class ConsoleApi:
  def __init__(self, window):
    window.keypad(True)
    window.scrollok(True)
    self.__window = window

  def __save_position(method):
    def wrapper(self, *args, **keyword_args):
      position = self.__window.getyx()
      method(self, *args, **keyword_args)
      self.__window.move(*position)
    return wrapper

  @property
  def screen_height(self):
    return self.__window.getmaxyx()[0]

  @property
  def screen_width(self):
    return self.__window.getmaxyx()[1]

  @property
  def x(self):
    return self.__window.getyx()[1]

  def get_char(self) -> str:
    return chr(self.__window.getch())

  @__save_position
  def put_char(self, char: str, attr=curses.A_NORMAL):
    assert self.__is_single_width_ascii_char(char)
    self.__window.addch(ord(char), attr)

  @__save_position
  def put_line(self, line_pos, text):
    assert len(text) <= self.screen_width
    self.__window.move(line_pos, 0)
    self.__window.clrtoeol()
    self.__window.addstr(line_pos, 0, text)

  def move(self, y, x):
    self.__window.move(y, x)

  def move_right(self):
    y, x = self.__window.getyx()
    if x == self.screen_width - 1: return
    self.__window.move(y, x + 1)

  def move_left(self):
    y, x = self.__window.getyx()
    if x == 0: return
    self.__window.move(y, x - 1)

  @__save_position
  def clear(self):
    self.__window.clear()

  @__save_position
  def scroll(self):
    self.__window.scroll()

  @staticmethod
  def __is_single_width_ascii_char(char: str):
    return len(char) == 1 and (curses.ascii.isprint(char) or char == " ")

class TypingGame:
  def __init__(self, api, example_file):
    self.__api = api
    self.__example_text = FormattedText(example_file,
                                        line_length=(api.screen_width - 1))
    self.__geometry = Geometry(api.screen_height)
    self.__input_text = InputText()
    if self.__example_text[0] == None:
      raise Exception("No line can be read from example source.")

  def play(self):
    self.__initialize_screen()

    game_over = False
    while not game_over:
      char = self.__api.get_char()

      if char in QUIT_CHARS:
        return
      elif char == CLEAR_CHAR:
        self.__clear_input_text()
      elif char in DELETE_CHARS:
        self.__delete_char()
      elif char == CHEAT_CHAR and CAN_CHEAT:
        game_over = self.__cheat()
      elif len(char) == 1 and (curses.ascii.isprint(char)
           or curses.ascii.isspace(char)):
        game_over = self.__add_char(char)

  def __add_char(self, char: str) -> bool:
    go_to_next_line = char == '\n' \
                      and self.__input_text == self.__example_text[0]
    if go_to_next_line and self.__example_text[1] == None:
      return True

    if go_to_next_line:
      self.__scroll()
    elif len(self.__input_text) == len(self.__example_text[0]):
      pass
    elif (char in string.printable
        and not unicodedata.category(char).startswith("C")) or char == "\t":
      self.__input_text.push_char(char)
      for printed_char in normalize_text(char):
        self.__print_char(printed_char)
    return False

  def __print_char(self, char: str):
    attr = ATTR_CORRECT if char == self.__example_text[0][self.__api.x] \
           else ATTR_ERROR
    self.__api.put_char(char, attr=attr)
    self.__api.move_right()

  def __delete_char(self):
    if len(self.__input_text) == 0: return
    for _ in range(len(normalize_text(self.__input_text.pop_char()))):
      self.__api.move_left()
      self.__api.put_char(self.__example_text[0][self.__api.x])

  def __cheat(self) -> bool:
    if self.__example_text[1] == None:
      return True
    self.__scroll()
    return False

  def __initialize_screen(self):
    self.__api.clear()
    self.__print_all_example_text()
    self.__clear_input_text()

  def __scroll(self):
    self.__api.scroll()
    del self.__example_text[0]
    self.__print_bottom_example_text()
    self.__clear_input_text()

  def __clear_input_text(self):
    self.__input_text = InputText()
    self.__api.put_line(self.__geometry.current_line, self.__example_text[0])
    self.__api.move(self.__geometry.current_line, 0)

  def __print_bottom_example_text(self):
    index = self.__geometry.bottom_line - self.__geometry.current_line
    if self.__example_text[index] != None:
      self.__api.put_line(self.__geometry.bottom_line,
                          self.__example_text[index])

  def __print_all_example_text(self):
    for index in range(self.__geometry.bottom_line
                       - self.__geometry.current_line + 1):
      if self.__example_text[index] == None: break
      self.__api.put_line(self.__geometry.current_line + index,
                          self.__example_text[index])

class InputText:
  def __init__(self):
    self.__text = ""

  def __eq__(self, text):
    return normalize_text(self.__text) == text

  def __len__(self):
    return len(normalize_text(self.__text))

  def push_char(self, char: str):
    self.__text += char

  def pop_char(self):
    last_char = self.__text[-1]
    self.__text = self.__text[:-1]
    return last_char

class Geometry:
  def __init__(self, screen_height):
    self.current_line = (screen_height - 1) // 2
    self.bottom_line = screen_height - 1

class FormattedText:
  def __init__(self, text_file, line_length=79):
    self.__file = text_file
    self.__line_length = line_length
    self.__buffered_lines = []

  def __del__(self):
    self.__file.close()

  def __getitem__(self, index: int) -> str:
    self.__read_lines_from_file()

    return self.__buffered_lines[index] if index < len(self.__buffered_lines) \
           else None

  def __delitem__(self, index: int):
    del self.__buffered_lines[index]

  def __read_lines_from_file(self):
    lines = self.__file.readlines()

    for line in lines:
      self.__buffer_line(normalize_text(line))

  def __buffer_line(self, line):
    line = line.rstrip()

    while len(line) > self.__line_length:
      buffered_line, line = self.__split_line(line)
      self.__buffered_lines.append(buffered_line)
    self.__buffered_lines.append(line)

  def __split_line(self, line):
    return line[:self.__line_length].rstrip(), \
           line[self.__line_length:].rstrip()

# functions

def normalize_text(text):
  return "".join(char for char in text_unidecode.unidecode(
                 text.replace("\t", " " * SPACES_PER_TAB))
                 if char in string.printable)

def reset_stdin():
  TEMPORARY_FD = 3

  os.dup2(sys.stdin.fileno(), TEMPORARY_FD)
  os.close(sys.stdin.fileno())
  sys.stdin = open(TTY_DEVICE_FILE)

  return os.fdopen(TEMPORARY_FD)

def initialize_curses():
  window = curses.initscr()
  curses.noecho()
  curses.cbreak()
  curses.start_color()
  curses.use_default_colors()
  return window

def finalize_curses():
  curses.nocbreak()
  curses.echo()
  curses.endwin() # should be at the very last line

def parse_args():
  arg_parser = argparse.ArgumentParser(description=DESCRIPTION)
  arg_parser.add_argument("example_path", nargs='?', default=None,
                          help="path or URI to example")
  arg_parser.add_argument("-c", "--cheat",
                          dest="can_cheat",
                          action="store_true",
                          help="enable the cheat key, {}"
                               .format(curses.ascii.unctrl(CHEAT_CHAR)))
  arg_parser.add_argument("-t", "--spaces-per-tab",
                          type=int, dest="spaces_per_tab",
                          help="set number of spaces per tab")
  arg_parser.add_argument("-v", "--version",
                          dest="show_version",
                          action="store_true", default=False,
                          help="show version information")

  args = arg_parser.parse_args()

  if args.show_version:
    print("version:", __version__)
    exit()

  if args.spaces_per_tab != None:
    global SPACES_PER_TAB
    SPACES_PER_TAB = args.spaces_per_tab
  if args.can_cheat == True:
    global CAN_CHEAT
    CAN_CHEAT = True

  return args

def get_example_file():
  example_path = parse_args().example_path

  if example_path != None and validators.url(example_path) \
      and sys.stdin.isatty():
    return get_remote_file(example_path)
  elif example_path != None and sys.stdin.isatty():
    return open(example_path)
  elif example_path == None and not sys.stdin.isatty():
    return reset_stdin()
  else:
    fail("Example is read from either the path or URI specified "
         "in the argument or stdin.")

def get_remote_file(uri):
  if urllib.parse.urlparse(uri).scheme not in SUPPORTED_SCHEMES:
    fail("Invalid scheme of URI is detected. "
         "(supported schemes: {})".format(", ".join(SUPPORTED_SCHEMES)))

  message("Loading page...")
  temporary_file = tempfile.TemporaryFile(mode="w+")
  with urllib.request.urlopen(uri) as response:
    temporary_file.write(response.read().decode(UTF8, "replace"))
  temporary_file.seek(0)
  return temporary_file

# main routine

def main():
  if not sys.stdout.isatty(): fail("stdout is not a tty.")

  example_file = get_example_file()

  try:
    # CAUTION:
    # You need to raise some Exception() instead of calling exit() here
    # to prevent curses from messing up your terminal.

    window = initialize_curses()

    TypingGame(ConsoleApi(window), example_file).play()
  finally:
    finalize_curses()

if __name__ == "__main__":
  main()
<|endoftext|>


<commit_before>"""
Milestones app initialization module
"""

__version__ = '0.3.0'
<commit_msg>[chore] Updating init file for 0.3.1 release<commit_after>"""
Milestones app initialization module
"""

__version__ = '0.3.1'
<|endoftext|>


<commit_before>
"""
Bubbles
=======

Locate bubbles in the system.

The following parameters apply to this filter:

.. glossary::

    Bubble species
        The species of the bubbles.

"""
import copy

import numpy as np

from . import base
from . import pointDefectsFilter
from . import speciesFilter
from . import clusterFilter

class BubblesFilterSettings(base.BaseSettings):
    """
    Setting for the bubbles filter.

    """
    def __init__(self):
        super(BubblesFilterSettings, self).__init__()

        self.registerSetting("bubbleSpecies", [])
        self.registerSetting("vacancyRadius", 1.3)
        self.registerSetting("vacNebRad", 4.0)

class BubblesFilter(base.BaseFilter):
    """
    The bubbles filter.

    """
    def apply(self, filterInput, settings):
        """Apply the bubbles filter."""
        if not isinstance(settings, BubblesFilterSettings):
            raise TypeError("BubbleFilter requires a settings object of type BubbleFilterSettings.")

        if not isinstance(filterInput, base.FilterInput):
            raise TypeError("BubbleFilter requires an input object of type FilterInput")

        # unpack input object
        visibleAtoms = filterInput.visibleAtoms
        inputState = filterInput.inputState
        refState = filterInput.refState

        # settings
        bubbleSpecies = settings.getSetting("bubbleSpecies")
        if not len(bubbleSpecies):
            self.logger.warning("No bubble species have been specified therefore no bubbles can be detected!")
            visibleAtoms.resize(0, refcheck=False)
            result = base.FilterResult()
            return result
        self.logger.debug("Bubble species: %r", bubbleSpecies)

        # clusters of bubble species and coms?

        # create new lattice with bubble species removed
        self.logger.debug("%d atoms in input before removing bubble species", inputState.NAtoms)
        lattice = copy.deepcopy(inputState)
        i = 0
        while i < lattice.NAtoms:
            if lattice.atomSym(i) in bubbleSpecies:
                lattice.removeAtom(i)
            else:
                i += 1
        self.logger.debug("%d atoms in input after removing bubble species", lattice.NAtoms)
        lattice.cellDims = lattice.cellDims9 # debugging
        lattice.writeLattice("noBubbleSpecies.dat")
        lattice.cellDims = lattice.cellDims3

        # locate clusters of vacancies
        defectsFilter = pointDefectsFilter.PointDefectsFilter("Point defects")
        defectsSettings = pointDefectsFilter.PointDefectsFilterSettings()
        defectsSettings.updateSetting("vacancyRadius", settings.getSetting("vacancyRadius"))
        defectsSettings.updateSetting("showInterstitials", False)
        defectsSettings.updateSetting("showAntisites", False)
        defectsSettings.updateSetting("findClusters", True)
        defectsSettings.updateSetting("neighbourRadius", settings.getSetting("vacNebRad"))
        defectsSettings.updateSetting("minClusterSize", 1)
        defectsSettings.updateSetting("visibleSpeciesList", lattice.specieList)

        defectsInput = base.FilterInput()
        defectsInput.inputState = lattice
        defectsInput.refState = refState
        defectsInput.interstitials = np.empty(lattice.NAtoms, np.int32)
        defectsInput.vacancies = np.empty(refState.NAtoms, np.int32)
        defectsInput.antisites = np.empty(refState.NAtoms, np.int32)
        defectsInput.onAntisites = np.empty(refState.NAtoms, np.int32)
        defectsInput.splitInterstitials = np.empty(3 * refState.NAtoms, np.int32)

        defectsResult = defectsFilter.apply(defectsInput, defectsSettings)
        vacClusters = defectsResult.getClusterList()
        self.logger.debug("%d vacancy clusters", len(vacClusters))

<commit_msg>Working on bubbles filter.<commit_after>
"""
Bubbles
=======

Locate bubbles in the system.

The following parameters apply to this filter:

.. glossary::

    Bubble species
        The species of the bubbles.

"""
import copy

import numpy as np

from . import base
from . import pointDefectsFilter
from . import speciesFilter
from . import clusterFilter

class BubblesFilterSettings(base.BaseSettings):
    """
    Setting for the bubbles filter.

    """
    def __init__(self):
        super(BubblesFilterSettings, self).__init__()

        self.registerSetting("bubbleSpecies", [])
        self.registerSetting("vacancyRadius", 1.3)
        self.registerSetting("vacNebRad", 4.0)

class BubblesFilter(base.BaseFilter):
    """
    The bubbles filter.

    """
    def apply(self, filterInput, settings):
        """Apply the bubbles filter."""
        if not isinstance(settings, BubblesFilterSettings):
            raise TypeError("BubbleFilter requires a settings object of type BubbleFilterSettings.")

        if not isinstance(filterInput, base.FilterInput):
            raise TypeError("BubbleFilter requires an input object of type FilterInput")

        # unpack input object
        visibleAtoms = filterInput.visibleAtoms
        inputState = filterInput.inputState
        refState = filterInput.refState

        # settings
        bubbleSpecies = settings.getSetting("bubbleSpecies")
        if not len(bubbleSpecies):
            self.logger.warning("No bubble species have been specified therefore no bubbles can be detected!")
            visibleAtoms.resize(0, refcheck=False)
            result = base.FilterResult()
            return result
        self.logger.debug("Bubble species: %r", bubbleSpecies)

        # clusters of bubble species and coms?

        # create new lattice with bubble species removed
        self.logger.debug("%d atoms in input before removing bubble species", inputState.NAtoms)
        lattice = copy.deepcopy(inputState)
        i = 0
        while i < lattice.NAtoms:
            if lattice.atomSym(i) in bubbleSpecies:
                lattice.removeAtom(i)
            else:
                i += 1
        self.logger.debug("%d atoms in input after removing bubble species", lattice.NAtoms)
        with open("noBubbleSpecies.dat", "w") as f:
            f.write("%d\n" % lattice.NAtoms)
            f.write("%f %f %f\n" % tuple(lattice.cellDims))
            for i in xrange(lattice.NAtoms):
                f.write("%s %f %f %f %f\n" % (lattice.atomSym(i), lattice.pos[3*i], lattice.pos[3*i+1], lattice.pos[3*i+2],
                                              lattice.charge[i]))

        # locate clusters of vacancies
        defectsFilter = pointDefectsFilter.PointDefectsFilter("Point defects")
        defectsSettings = pointDefectsFilter.PointDefectsFilterSettings()
        defectsSettings.updateSetting("vacancyRadius", settings.getSetting("vacancyRadius"))
        defectsSettings.updateSetting("showInterstitials", False)
        defectsSettings.updateSetting("showAntisites", False)
        defectsSettings.updateSetting("findClusters", True)
        defectsSettings.updateSetting("neighbourRadius", settings.getSetting("vacNebRad"))
        defectsSettings.updateSetting("minClusterSize", 1)
        defectsSettings.updateSetting("visibleSpeciesList", lattice.specieList)

        defectsInput = base.FilterInput()
        defectsInput.inputState = lattice
        defectsInput.refState = refState
        defectsInput.interstitials = np.empty(lattice.NAtoms, np.int32)
        defectsInput.vacancies = np.empty(refState.NAtoms, np.int32)
        defectsInput.antisites = np.empty(refState.NAtoms, np.int32)
        defectsInput.onAntisites = np.empty(refState.NAtoms, np.int32)
        defectsInput.splitInterstitials = np.empty(3 * refState.NAtoms, np.int32)

        defectsResult = defectsFilter.apply(defectsInput, defectsSettings)
        vacClusters = defectsResult.getClusterList()
        self.logger.debug("%d vacancy clusters", len(vacClusters))

<|endoftext|>


<commit_before>"""
minfraud.webservice
~~~~~~~~~~~~~~~~~~~

This module contains the webservice client class.

"""

import requests
from requests.utils import default_user_agent
from voluptuous import MultipleInvalid

from .version import __version__
from .errors import MinFraudError, HTTPError, AuthenticationError, \
    InsufficientFundsError, InvalidRequestError, PermissionRequiredError
from .models import Factors, Insights, Score
from .validation import validate_transaction

class Client(object):
    """Client for accessing the minFraud Score and Insights web services."""

    def __init__(self,
                 user_id,
                 license_key,
                 host='minfraud.maxmind.com',
                 locales=('en', ),
                 timeout=None):
        """Constructor for Client.

        :param user_id: Your MaxMind user ID
        :type user_id: int
        :param license_key: Your MaxMind license key
        :type license_key: str
        :param host: The host to use when connecting to the web service.
        :type host: str
        :param locales: A tuple of locale codes to use in name property
        :type locales: tuple[str]
        :param timeout: The timeout to use for the request.
        :type timeout: float
        :return: Client object
        :rtype: Client
        """
        # pylint: disable=too-many-arguments
        self._locales = locales
        # requests 2.12.2 requires that the username passed to auth be a
        # string
        self._user_id = user_id if isinstance(user_id, bytes) else str(user_id)
        self._license_key = license_key
        self._base_uri = u'https://{0:s}/minfraud/v2.0'.format(host)
        self._timeout = timeout

    def factors(self, transaction, validate=True):
        """Query Factors endpoint with transaction data.

        :param transaction: A dictionary containing the transaction to be
          sent to the minFraud Insights web service as specified in the `REST
          API documentation
          <https://dev.maxmind.com/minfraud/#Request_Body>`_.
        :type transaction: dict
        :param validate: If set to false, validation of the transaction
          dictionary will be disabled. This validation helps ensure that your
          request is correct before sending it to MaxMind. Validation raises an
          InvalidRequestError.
        :type validate: bool
        :return: A Factors model object
        :rtype: Factors
        :raises: AuthenticationError, InsufficientFundsError,
          InvalidRequestError, HTTPError, MinFraudError,
        """
        return self._response_for('factors', Factors, transaction, validate)

    def insights(self, transaction, validate=True):
        """Query Insights endpoint with transaction data.

        :param transaction: A dictionary containing the transaction to be
          sent to the minFraud Insights web service as specified in the `REST
          API documentation
          <https://dev.maxmind.com/minfraud/#Request_Body>`_.
        :type transaction: dict
        :param validate: If set to false, validation of the transaction
          dictionary will be disabled. This validation helps ensure that your
          request is correct before sending it to MaxMind. Validation raises an
          InvalidRequestError.
        :type validate: bool
        :return: An Insights model object
        :rtype: Insights
        :raises: AuthenticationError, InsufficientFundsError,
          InvalidRequestError, HTTPError, MinFraudError,
        """
        return self._response_for('insights', Insights, transaction, validate)

    def score(self, transaction, validate=True):
        """Query Score endpoint with transaction data.

        :param transaction: A dictionary containing the transaction to be
          sent to the minFraud Score web service as specified in the `REST API
          documentation
          <https://dev.maxmind.com/minfraud/#Request_Body>`_.
        :type transaction: dict
        :param validate: If set to false, validation of the transaction
          dictionary will be disabled. This validation helps ensure that your
          request is correct before sending it to MaxMind. Validation raises an
          InvalidRequestError.
        :type validate: bool
        :return: A Score model object
        :rtype: Score
        :raises: AuthenticationError, InsufficientFundsError,
          InvalidRequestError, HTTPError, MinFraudError,
        """
        return self._response_for('score', Score, transaction, validate)

    def _response_for(self, path, model_class, request, validate):
        """Send request and create response object."""
        cleaned_request = self._copy_and_clean(request)
        if validate:
            try:
                validate_transaction(cleaned_request)
            except MultipleInvalid as ex:
                raise InvalidRequestError(
                    "Invalid transaction data: {0}".format(ex))
        uri = '/'.join([self._base_uri, path])
        response = requests.post(
            uri,
            json=cleaned_request,
            auth=(self._user_id, self._license_key),
            headers={
                'Accept': 'application/json',
                'User-Agent': self._user_agent()
            },
            timeout=self._timeout)
        if response.status_code == 200:
            return self._handle_success(response, uri, model_class)
        else:
            self._handle_error(response, uri)

    def _copy_and_clean(self, data):
        """Create a copy of the data structure with Nones removed."""
        if isinstance(data, dict):
            return dict((k, self._copy_and_clean(v))
                        for (k, v) in data.items() if v is not None)
        elif isinstance(data, (list, set, tuple)):
            return [self._copy_and_clean(x) for x in data if x is not None]
        return data

    def _user_agent(self):
        """Create User-Agent header."""
        return 'minFraud-API/%s %s' % (__version__, default_user_agent())

    def _handle_success(self, response, uri, model_class):
        """Handle successful response."""
        try:
            body = response.json()
        except ValueError:
            raise MinFraudError('Received a 200 response'
                                ' but could not decode the response as '
                                'JSON: {0}'.format(response.content), 200, uri)
        if 'ip_address' in body:
            body['ip_address']['_locales'] = self._locales
        return model_class(body)

    def _handle_error(self, response, uri):
        """Handle error responses."""
        status = response.status_code

        if 400 <= status < 500:
            self._handle_4xx_status(response, status, uri)
        elif 500 <= status < 600:
            self._handle_5xx_status(status, uri)
        else:
            self._handle_non_200_status(status, uri)

    def _handle_4xx_status(self, response, status, uri):
        """Handle error responses with 4xx status codes."""
        if not response.content:
            raise HTTPError('Received a {0} error with no body'.format(status),
                            status, uri)
        elif response.headers.get('Content-Type', '').find('json') == -1:
            raise HTTPError('Received a {0} with the following '
                            'body: {1}'.format(status, response.content),
                            status, uri)
        try:
            body = response.json()
        except ValueError:
            raise HTTPError(
                'Received a {status:d} error but it did not include'
                ' the expected JSON body: {content}'.format(
                    status=status, content=response.content), status, uri)
        else:
            if 'code' in body and 'error' in body:
                self._handle_web_service_error(
                    body.get('error'), body.get('code'), status, uri)
            else:
                raise HTTPError(
                    'Error response contains JSON but it does not specify code'
                    ' or error keys: {0}'.format(response.content), status,
                    uri)

    def _handle_web_service_error(self, message, code, status, uri):
        """Handle error responses with the JSON body from the web service."""
        if code in ('AUTHORIZATION_INVALID', 'LICENSE_KEY_REQUIRED',
                    'USER_ID_REQUIRED'):
            raise AuthenticationError(message)
        elif code == 'INSUFFICIENT_FUNDS':
            raise InsufficientFundsError(message)
        elif code == 'PERMISSION_REQUIRED':
            raise PermissionRequiredError(message)

        raise InvalidRequestError(message, code, status, uri)

    def _handle_5xx_status(self, status, uri):
        """Handle error response with 5xx status codes."""
        raise HTTPError(u'Received a server error ({0}) for '
                        u'{1}'.format(status, uri), status, uri)

    def _handle_non_200_status(self, status, uri):
        """Handle successful responses with unexpected status codes."""
        raise HTTPError(u'Received an unexpected HTTP status '
                        u'({0}) for {1}'.format(status, uri), status, uri)
<commit_msg>Change error handler code to return exception<commit_after>"""
minfraud.webservice
~~~~~~~~~~~~~~~~~~~

This module contains the webservice client class.

"""

import requests
from requests.utils import default_user_agent
from voluptuous import MultipleInvalid

from .version import __version__
from .errors import MinFraudError, HTTPError, AuthenticationError, \
    InsufficientFundsError, InvalidRequestError, PermissionRequiredError
from .models import Factors, Insights, Score
from .validation import validate_transaction

class Client(object):
    """Client for accessing the minFraud Score and Insights web services."""

    def __init__(self,
                 user_id,
                 license_key,
                 host='minfraud.maxmind.com',
                 locales=('en', ),
                 timeout=None):
        """Constructor for Client.

        :param user_id: Your MaxMind user ID
        :type user_id: int
        :param license_key: Your MaxMind license key
        :type license_key: str
        :param host: The host to use when connecting to the web service.
        :type host: str
        :param locales: A tuple of locale codes to use in name property
        :type locales: tuple[str]
        :param timeout: The timeout to use for the request.
        :type timeout: float
        :return: Client object
        :rtype: Client
        """
        # pylint: disable=too-many-arguments
        self._locales = locales
        # requests 2.12.2 requires that the username passed to auth be a
        # string
        self._user_id = user_id if isinstance(user_id, bytes) else str(user_id)
        self._license_key = license_key
        self._base_uri = u'https://{0:s}/minfraud/v2.0'.format(host)
        self._timeout = timeout

    def factors(self, transaction, validate=True):
        """Query Factors endpoint with transaction data.

        :param transaction: A dictionary containing the transaction to be
          sent to the minFraud Insights web service as specified in the `REST
          API documentation
          <https://dev.maxmind.com/minfraud/#Request_Body>`_.
        :type transaction: dict
        :param validate: If set to false, validation of the transaction
          dictionary will be disabled. This validation helps ensure that your
          request is correct before sending it to MaxMind. Validation raises an
          InvalidRequestError.
        :type validate: bool
        :return: A Factors model object
        :rtype: Factors
        :raises: AuthenticationError, InsufficientFundsError,
          InvalidRequestError, HTTPError, MinFraudError,
        """
        return self._response_for('factors', Factors, transaction, validate)

    def insights(self, transaction, validate=True):
        """Query Insights endpoint with transaction data.

        :param transaction: A dictionary containing the transaction to be
          sent to the minFraud Insights web service as specified in the `REST
          API documentation
          <https://dev.maxmind.com/minfraud/#Request_Body>`_.
        :type transaction: dict
        :param validate: If set to false, validation of the transaction
          dictionary will be disabled. This validation helps ensure that your
          request is correct before sending it to MaxMind. Validation raises an
          InvalidRequestError.
        :type validate: bool
        :return: An Insights model object
        :rtype: Insights
        :raises: AuthenticationError, InsufficientFundsError,
          InvalidRequestError, HTTPError, MinFraudError,
        """
        return self._response_for('insights', Insights, transaction, validate)

    def score(self, transaction, validate=True):
        """Query Score endpoint with transaction data.

        :param transaction: A dictionary containing the transaction to be
          sent to the minFraud Score web service as specified in the `REST API
          documentation
          <https://dev.maxmind.com/minfraud/#Request_Body>`_.
        :type transaction: dict
        :param validate: If set to false, validation of the transaction
          dictionary will be disabled. This validation helps ensure that your
          request is correct before sending it to MaxMind. Validation raises an
          InvalidRequestError.
        :type validate: bool
        :return: A Score model object
        :rtype: Score
        :raises: AuthenticationError, InsufficientFundsError,
          InvalidRequestError, HTTPError, MinFraudError,
        """
        return self._response_for('score', Score, transaction, validate)

    def _response_for(self, path, model_class, request, validate):
        """Send request and create response object."""
        cleaned_request = self._copy_and_clean(request)
        if validate:
            try:
                validate_transaction(cleaned_request)
            except MultipleInvalid as ex:
                raise InvalidRequestError(
                    "Invalid transaction data: {0}".format(ex))
        uri = '/'.join([self._base_uri, path])
        response = requests.post(
            uri,
            json=cleaned_request,
            auth=(self._user_id, self._license_key),
            headers={
                'Accept': 'application/json',
                'User-Agent': self._user_agent()
            },
            timeout=self._timeout)
        if response.status_code != 200:
            raise self._exception_for_error(response, uri)
        return self._handle_success(response, uri, model_class)

    def _copy_and_clean(self, data):
        """Create a copy of the data structure with Nones removed."""
        if isinstance(data, dict):
            return dict((k, self._copy_and_clean(v))
                        for (k, v) in data.items() if v is not None)
        elif isinstance(data, (list, set, tuple)):
            return [self._copy_and_clean(x) for x in data if x is not None]
        return data

    def _user_agent(self):
        """Create User-Agent header."""
        return 'minFraud-API/%s %s' % (__version__, default_user_agent())

    def _handle_success(self, response, uri, model_class):
        """Handle successful response."""
        try:
            body = response.json()
        except ValueError:
            raise MinFraudError('Received a 200 response'
                                ' but could not decode the response as '
                                'JSON: {0}'.format(response.content), 200, uri)
        if 'ip_address' in body:
            body['ip_address']['_locales'] = self._locales
        return model_class(body)

    def _exception_for_error(self, response, uri):
        """Returns the exception for the error responses."""
        status = response.status_code

        if 400 <= status < 500:
            return self._exception_for_4xx_status(response, status, uri)
        elif 500 <= status < 600:
            return self._exception_for_5xx_status(status, uri)
        return self._exception_for_non_200_status(status, uri)

    def _exception_for_4xx_status(self, response, status, uri):
        """Returns exception for error responses with 4xx status codes."""
        if not response.content:
            return HTTPError(
                'Received a {0} error with no body'.format(status), status,
                uri)
        elif response.headers.get('Content-Type', '').find('json') == -1:
            return HTTPError('Received a {0} with the following '
                             'body: {1}'.format(status, response.content),
                             status, uri)
        try:
            body = response.json()
        except ValueError:
            return HTTPError(
                'Received a {status:d} error but it did not include'
                ' the expected JSON body: {content}'.format(
                    status=status, content=response.content), status, uri)
        else:
            if 'code' in body and 'error' in body:
                return self._exception_for_web_service_error(
                    body.get('error'), body.get('code'), status, uri)
            return HTTPError(
                'Error response contains JSON but it does not specify code'
                ' or error keys: {0}'.format(response.content), status, uri)

    def _exception_for_web_service_error(self, message, code, status, uri):
        """Returns exception for error responses with the JSON body."""
        if code in ('AUTHORIZATION_INVALID', 'LICENSE_KEY_REQUIRED',
                    'USER_ID_REQUIRED'):
            return AuthenticationError(message)
        elif code == 'INSUFFICIENT_FUNDS':
            return InsufficientFundsError(message)
        elif code == 'PERMISSION_REQUIRED':
            return PermissionRequiredError(message)

        return InvalidRequestError(message, code, status, uri)

    def _exception_for_5xx_status(self, status, uri):
        """Returns exception for error response with 5xx status codes."""
        return HTTPError(u'Received a server error ({0}) for '
                         u'{1}'.format(status, uri), status, uri)

    def _exception_for_non_200_status(self, status, uri):
        """Returns exception for responses with unexpected status codes."""
        return HTTPError(u'Received an unexpected HTTP status '
                         u'({0}) for {1}'.format(status, uri), status, uri)
<|endoftext|>


<commit_before>from numpy.testing import assert_array_equal
from nose.tools import (
    assert_equal,
    assert_almost_equal,
    assert_dict_equal,
    assert_list_equal,
    assert_tuple_equal,
    assert_raises,
    )

def check_equal(a, b, msg=None):
    assert_equal(a, b, msg)

def check_list_equal(a, b, msg=None):
    assert_list_equal(a, b, msg)

def check_dict_equal(a, b, msg=None):
    assert_dict_equal(a, b, msg)

def check_array_equal(a, b, **kwargs):
    assert_array_equal(a, b, **kwargs)

def check_tuple_equal(a, b, msg=None):
    assert_tuple_equal(a, b, msg)

def check_raises(cls, callable_, *args):
    assert_raises(cls, callable_, *args)

<commit_msg>more test methods<commit_after>from numpy.testing import (
    assert_array_equal,
    assert_almost_equal,
)

from nose.tools import (
    assert_equal,
    assert_true,
#    assert_almost_equal,
    assert_dict_equal,
    assert_list_equal,
    assert_tuple_equal,
    assert_raises,
    )

def check_equal(a, b, msg=None):
    assert_equal(a, b, msg)

def check_array_equal(a, b, **kwargs):
    assert_array_equal(a, b, **kwargs)

def check_almost_equal(a,b,kwargs={}):
    assert_almost_equal(a,b, **kwargs)

def check_true(a, kwargs={}):
    assert_true(a, **kwargs)

def check_not_equal(a, b):
    assert_raises(AssertionError, check_array_equal, a, b)

def check_list_equal(a, b, msg=None):
    assert_list_equal(a, b, msg)

def check_dict_equal(a, b, msg=None):
    assert_dict_equal(a, b, msg)

def check_tuple_equal(a, b, msg=None):
    assert_tuple_equal(a, b, msg)

def check_raises(cls, callable_, *args):
    assert_raises(cls, callable_, *args)

<|endoftext|>


<commit_before># Copyright 2015 Huawei Technologies Co.,LTD.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc

from oslo_config import cfg
from oslo_log import log
import six

from magnum import objects
from magnum.objects.fields import BayType as bay_type

LOG = log.getLogger(__name__)

CONF = cfg.CONF
CONF.import_opt('docker_remote_api_version',
                'magnum.conductor.handlers.docker_conductor',
                group='docker')
CONF.import_opt('default_timeout',
                'magnum.conductor.handlers.docker_conductor',
                group='docker')

@six.add_metaclass(abc.ABCMeta)
class MonitorBase(object):

    def __init__(self, context, bay):
        self.context = context
        self.bay = bay

    @abc.abstractproperty
    def metrics_spec(self):
        """Metric specification."""

    @abc.abstractmethod
    def pull_data(self):
        """Pull data for monitoring."""

    def get_metric_names(self):
        return self.metrics_spec.keys()

    def get_metric_unit(self, metric_name):
        return self.metrics_spec[metric_name]['unit']

    def compute_metric_value(self, metric_name):
        func_name = self.metrics_spec[metric_name]['func']
        func = getattr(self, func_name)
        return func()

def create_monitor(context, bay):
    baymodel = objects.BayModel.get_by_uuid(context, bay.baymodel_id)
    if baymodel.coe == bay_type.SWARM:
        from magnum.conductor.swarm_monitor import SwarmMonitor
        return SwarmMonitor(context, bay)
    elif baymodel.coe == bay_type.KUBERNETES:
        from magnum.conductor.k8s_monitor import K8sMonitor
        return K8sMonitor(context, bay)

    # TODO(hongbin): add support for other bay types
    LOG.debug("Cannot create monitor with bay type '%s'" % baymodel.coe)
    return None
<commit_msg>use importutils in monitors.py to avoid cyclic imports<commit_after># Copyright 2015 Huawei Technologies Co.,LTD.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc

from oslo_config import cfg
from oslo_log import log
from oslo_utils import importutils
import six

from magnum import objects
from magnum.objects.fields import BayType as bay_type

LOG = log.getLogger(__name__)

CONF = cfg.CONF
CONF.import_opt('docker_remote_api_version',
                'magnum.conductor.handlers.docker_conductor',
                group='docker')
CONF.import_opt('default_timeout',
                'magnum.conductor.handlers.docker_conductor',
                group='docker')

COE_CLASS_PATH = {
    bay_type.SWARM: 'magnum.conductor.swarm_monitor.SwarmMonitor',
    bay_type.KUBERNETES: 'magnum.conductor.k8s_monitor.K8sMonitor'
}

@six.add_metaclass(abc.ABCMeta)
class MonitorBase(object):

    def __init__(self, context, bay):
        self.context = context
        self.bay = bay

    @abc.abstractproperty
    def metrics_spec(self):
        """Metric specification."""

    @abc.abstractmethod
    def pull_data(self):
        """Pull data for monitoring."""

    def get_metric_names(self):
        return self.metrics_spec.keys()

    def get_metric_unit(self, metric_name):
        return self.metrics_spec[metric_name]['unit']

    def compute_metric_value(self, metric_name):
        func_name = self.metrics_spec[metric_name]['func']
        func = getattr(self, func_name)
        return func()

def create_monitor(context, bay):
    baymodel = objects.BayModel.get_by_uuid(context, bay.baymodel_id)
    if baymodel.coe in COE_CLASS_PATH:
        coe_cls = importutils.import_class(COE_CLASS_PATH[baymodel.coe])
        return coe_cls(context, bay)

    # TODO(hongbin): add support for other bay types
    LOG.debug("Cannot create monitor with bay type '%s'" % baymodel.coe)
    return None
<|endoftext|>


<commit_before>"""
MagPy
MagPy's CDF input/output filters
Rewritten by Roman Leonhardt October 2019
- contains test, read and write methods
        MagPyCDF
- based on cdflib
- supports python >= 3.5

"""
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division
from io import open

# Specify what methods are really needed
from magpy.stream import *

import logging
logger = logging.getLogger(__name__)

import gc
import cdflib

def isPYCDF(filename):
    """
    Checks whether a file is Nasa CDF format.
    """
    try:
        temp = cdflib.CDF(filename)
    except:
        return False
    try:
        cdfformat = temp.globalattsget().get('DataFormat')
        if not cdfformat.startswith('MagPyCDF'):
            return False
    except:
        pass
    try:
        variables = temp.cdf_info().get('zVariables')
        if not 'Epoch' in variables:
            if not 'time' in variables:
                return False
    except:
        return False

    logger.debug("format_magpy: Found PYCDF file %s" % filename)
    return True

def readPYCDF(filename, headonly=False, **kwargs):
    """
    Reading CDF format data.
    """
    stream = DataStream([],{},np.asarray([[] for key in KEYLIST]))

    array = [[] for key in KEYLIST]
    starttime = kwargs.get('starttime')
    endtime = kwargs.get('endtime')
    oldtype = kwargs.get('oldtype')
    getfile = True

    print ("Reading PYCDF with CDFLIB")
    # Check whether header information is already present
    headskip = False
    if stream.header == None:
        stream.header.clear()
    else:
        headskip = True

    theday = extractDateFromString(filename)
    try:
        if starttime:
            if not theday[-1] >= datetime.date(stream._testtime(starttime)):
                getfile = False
        if endtime:
            if not theday[0] <= datetime.date(stream._testtime(endtime)):
                getfile = False
    except:
        # Date format not recognized. Need to read all files
        getfile = True
    logbaddata = False

    # Get format type:
    # MagPy type is using datetime objects
    if getfile:
        cdfdat = cdflib.CDF(filename)
        cdfformat = cdfdat.globalattsget().get('DataFormat')
        try:
            version = float(cdfformat.replace('MagPyCDF',''))
        except:
            version = 1.0

        if headskip:  # TODO find out why
            for att in cdfdat.globalattsget():
                value = cdfdat.globalattsget().get(att)
                try:
                    if isinstance(list(value), list):
                        if len(value) == 1:
                            value = value[0]
                except:
                    pass
                if not att in ['DataAbsFunctionObject','DataBaseValues', 'DataFlagList']:
                    stream.header[att] = value
                else:
                        print ("Found special header content !!!!!!!!!!!!!!!! --  version {}".format(version))
                        #TODO check this - is pickle really necessary?
                        logger.debug("readPYCDF: Found object - loading and unpickling")
                        func = ''
                        try:
                            func = pickle.loads(str.encode(value), encoding="bytes")
                        except:
                            try:
                                print ("old unpickling version")
                                func = pickle.loads(value)
                            except:
                                print ("FAILED to load special content")
                                logger.debug("readPYCDF: Failed to load Object - constructed before v0.2.000?")
                        stream.header[att] = func

        logger.info('readPYCDF: %s Format: %s ' % (filename, cdfformat))

        #if debug:
        #print ("Step2", stream.header)

        variables = cdfdat.cdf_info().get('zVariables')
        timelength = 0

        print (variables)
        for key in variables:
            if key.find('time')>=0 or key == 'Epoch':
                # Time column identified
                if not key == 'sectime':
                    ind = KEYLIST.index('time')
                else:
                    ind = KEYLIST.index('sectime')
                try:
                    ttdesc = cdfdat.varinq(key).get('Data_Type_Description')
                    if not ttdesc == 'CDF_TIME_TT2000':
                        print ("WARNING: Time column is not CDF_TIME_TT2000 (found {})".format(ttdesc))
                    col = cdfdat.varget(key)
                    array[ind] = date2num(np.asarray([datetime.utcfromtimestamp(el) for el in cdflib.cdfepoch.unixtime(col)]))
                except:
                    array[ind] = np.asarray([])
            else:
                ind = KEYLIST.index(key)
                addhead = False
                timelength = len(array[0])
                ttdesc = cdfdat.varinq(key).get('Data_Type_Description')
                col = cdfdat.varget(key)
                #print (ttdesc, cdfdat.varinq(key).get('Data_Type'))
                print (key, len(col), col)
                #print ([ele for ele in col if not ele ==''])
                if not len(col) == timelength and len(col) == 1:
                    array[ind] = np.asarray([col[0]]*timelength)
                    addhead = True
                elif not len(col) == timelength and len(col)>1:
                    array[ind] = np.asarray([])
                else:
                    array[ind] = np.asarray(col)
                    addhead = True
                #print (cdfdat.varattsget(key))
                if addhead:
                    vname = cdfdat.varattsget(key).get('name','')
                    if not vname:
                        vname = cdfdat.varattsget(key).get('FIELDNAM','')
                    vunit = cdfdat.varattsget(key).get('units','')
                    if not vunit:
                        vunit = cdfdat.varattsget(key).get('UNITS','')
                    stream.header['col-'+key.lower()] = vname
                    stream.header['unit-col-'+key.lower()] = vunit

        cdfdat.close()
        del cdfdat

    #print (stream.header)
    return DataStream([LineStruct()], stream.header,np.asarray(array))

def writePYCDF(datastream, filename, **kwargs):
    """
    VARIABLES
        new: use compression variable instead of skipcompression
        compression = 0: skip compression
        compression = 1-9: use this compression factor:
                           9 high compreesion (slow)
                           1 low compression (fast)
               default is 5

    """

    if pyvers and pyvers == 2:
                ch1 = '-'.encode('utf-8') # not working with py3
                ch2 = ''.encode('utf-8')
    else:
                ch1 = '-'
                ch2 = ''

    if not len(datastream.ndarray[0]) > 0 and not len(datastream) > 0:
        return False

    def tt(my_dt_ob):
        ms = my_dt_ob.microsecond/1000.  # fraction
        date_list = [my_dt_ob.year, my_dt_ob.month, my_dt_ob.day, my_dt_ob.hour, my_dt_ob.minute, my_dt_ob.second, ms]
        return date_list

    logger.info("Writing PYCDF Format {}".format(filename))
    mode = kwargs.get('mode')
    skipcompression = kwargs.get('skipcompression')
    compression = kwargs.get('compression')
    version = '1.2'

    if compression == 0: ## temporary solution until all refs to skipcomression are eliminated
        skipcompression = True

    main_cdf_spec = {}
    main_cdf_spec['Compressed'] = False

    leapsecondlastupdate = cdflib.cdfepoch.getLeapSecondLastUpdated()

    if not skipcompression:
        try:
            main_cdf_spec['Compressed'] = True
        except:
            logger.warning("writeIMAGCDF: Compression failed for unknown reason - storing uncompresed data")
            pass

    testname = str(filename+'.cdf')

    if os.path.isfile(testname):
        filename = testname
    if os.path.isfile(filename):
        if mode == 'skip': # skip existing inputs
            exst = read(path_or_url=filename)
            datastream = joinStreams(exst,datastream,extend=True)
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
        elif mode == 'replace': # replace existing inputs
            #print filename
            #### !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            ####  Please note: Replacing requires a lot memory
            #### If memory issues appear then please overwrite existing data
            #### TODO Optimze sorting
            #### !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            try:
                exst = read(path_or_url=filename)
                datastream = joinStreams(datastream,exst,extend=True)
            except:
                logger.error("writePYCDF: Could not interprete existing data set - aborting")
                sys.exit()
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
        elif mode == 'append':
            #print filename
            exst = read(path_or_url=filename)
            datastream = joinStreams(exst,datastream,extend=True)
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
        else: # overwrite mode
            #print filename
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
    else:
        mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)

    keylst = datastream._get_key_headers()

    if not 'flag' in keylst:
        keylst.append('flag')
    if not 'comment' in keylst:
        keylst.append('comment')
    if not 'typ' in keylst:
        keylst.append('typ')
    tmpkeylst = ['time']
    tmpkeylst.extend(keylst)
    keylst = tmpkeylst

    headdict = datastream.header
    head, line = [],[]
    globalAttrs = {}

    if not mode == 'append':
        for key in headdict:
            if not key.find('col-') >= 0:
                #print (key, headdict[key])
                if not key in ['DataAbsFunctionObject','DataBaseValues', 'DataFlagList']:
                    globalAttrs[key] = { 0 : str(headdict[key]) }
                else:
                    logger.info("writePYCDF: Found Object in header - pickle and dump ")
                    pfunc = pickle.dumps(headdict[key])
                    globalAttrs[key] = { 0 : pfunc }

    globalAttrs['DataFormat'] = { 0 : 'MagPyCDF{}'.format(version)}

    mycdf.write_globalattrs(globalAttrs)

    def checkEqual3(lst):
        return lst[1:] == lst[:-1]

    for key in keylst:
        var_attrs = {}
        var_spec = {}
        cdfdata = []

        ind = KEYLIST.index(key)
        col = datastream.ndarray[ind]

        if not key in NUMKEYLIST:
            #if not key == 'time':  ### why??
            col = np.asarray(col)

        # Sort out columns only containing nan's
        try:
            test = [elem for elem in col if not isnan(elem)]
            if not len(test) > 0:
                col = np.asarray([])
        except:
            pass

        if not False in checkEqual3(col) and len(col) > 0:
            logger.warning("writePYCDF: Found identical values only for key: %s" % key)
            col = col[:1]

        cdfkey = key.lower()
        if key.find('time') >= 0:
            if cdfkey.endswith('time'):
                if cdfkey == 'time':
                    cdfkey = 'Epoch'
                try: # Might fail for sectime
                    cdfdata = cdflib.cdfepoch.compute_tt2000( [tt(num2date(elem).replace(tzinfo=None)) for elem in col.astype(np.float64)] )
                except:
                    cdfdata = np.asarray([])
                var_spec['Data_Type'] = 33
        elif len(col) > 0:
            print (cdfkey)
            if not key in NUMKEYLIST:
                col = list(col)
                col = ['' if el is None else el for el in col]
                col = np.asarray(col) # to get string conversion
                col = list(col) # convert back to list for write_var
                var_spec['Data_Type'] = 51
            else:
                var_spec['Data_Type'] = 45
                col = np.asarray([np.nan if el in [None,ch1] else el for el in col])
                col = col.astype(float)
            cdfdata = col

            var_attrs['name'] = headdict.get('col-'+key,'')       # use 'FIELDNAM' to be conform with NASA / IMAGCDF style ## Version 1.1
            var_attrs['units'] = headdict.get('unit-col-'+key,'')  # use 'UNITS' to be conform with NASA style / IMAGCDF style
            var_attrs['FIELDNAM'] = headdict.get('col-'+key,'')       # use 'FIELDNAM' to be conform with NASA / IMAGCDF style ## Version 1.2
            var_attrs['UNITS'] = headdict.get('unit-col-'+key,'')  # use 'UNITS' to be conform with NASA style / IMAGCDF style ## Version 1.2
            var_attrs['LABLAXIS'] = headdict.get('col-'+key,'') ## Version 1.2

        #if cdfdata is just a single value, then it will be converted into an array
        try:
            test = len(cdfdata)
        except:
            print ("Failed")
            cdfdata = np.asarray([cdfdata])

        if len(cdfdata) > 0:
            var_spec['Variable'] = cdfkey
            var_spec['Num_Elements'] = 1
            var_spec['Rec_Vary'] = True # The dimensional sizes, applicable only to rVariables.
            var_spec['Dim_Sizes'] = []

            mycdf.write_var(var_spec, var_attrs=var_attrs, var_data=cdfdata)

    mycdf.close()
    return testname

<commit_msg>pycdf can now read and write flags in py3<commit_after>"""
MagPy
MagPy's CDF input/output filters
Rewritten by Roman Leonhardt October 2019
- contains test, read and write methods
        MagPyCDF
- based on cdflib
- supports python >= 3.5

"""
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division
from io import open

# Specify what methods are really needed
from magpy.stream import *

import logging
logger = logging.getLogger(__name__)

import gc
import cdflib

def isPYCDF(filename):
    """
    Checks whether a file is Nasa CDF format.
    """
    try:
        temp = cdflib.CDF(filename)
    except:
        return False
    try:
        cdfformat = temp.globalattsget().get('DataFormat')
        if not cdfformat.startswith('MagPyCDF'):
            return False
    except:
        pass
    try:
        variables = temp.cdf_info().get('zVariables')
        if not 'Epoch' in variables:
            if not 'time' in variables:
                return False
    except:
        return False

    logger.debug("format_magpy: Found PYCDF file %s" % filename)
    return True

def readPYCDF(filename, headonly=False, **kwargs):
    """
    Reading CDF format data.
    """
    stream = DataStream([],{},np.asarray([[] for key in KEYLIST]))

    array = [[] for key in KEYLIST]
    starttime = kwargs.get('starttime')
    endtime = kwargs.get('endtime')
    oldtype = kwargs.get('oldtype')
    getfile = True

    print ("Reading PYCDF with CDFLIB")
    # Check whether header information is already present
    headskip = False
    if stream.header == None:
        stream.header.clear()
    else:
        headskip = True

    theday = extractDateFromString(filename)
    try:
        if starttime:
            if not theday[-1] >= datetime.date(stream._testtime(starttime)):
                getfile = False
        if endtime:
            if not theday[0] <= datetime.date(stream._testtime(endtime)):
                getfile = False
    except:
        # Date format not recognized. Need to read all files
        getfile = True
    logbaddata = False

    # Get format type:
    # MagPy type is using datetime objects
    if getfile:
        cdfdat = cdflib.CDF(filename)
        cdfformat = cdfdat.globalattsget().get('DataFormat')
        try:
            version = float(cdfformat.replace('MagPyCDF',''))
        except:
            version = 1.0

        if headskip:  # TODO find out why
            for att in cdfdat.globalattsget():
                value = cdfdat.globalattsget().get(att)
                try:
                    if isinstance(list(value), list):
                        if len(value) == 1:
                            value = value[0]
                except:
                    pass
                if not att in ['DataAbsFunctionObject','DataBaseValues', 'DataFlagList']:
                    stream.header[att] = value
                else:
                        print ("Found special header content !!!!!!!!!!!!!!!! --  version {}".format(version))
                        #TODO check this - is pickle really necessary?
                        logger.debug("readPYCDF: Found object - loading and unpickling")
                        func = ''
                        try:
                            func = pickle.loads(str.encode(value), encoding="bytes")
                        except:
                            try:
                                print ("old unpickling version")
                                func = pickle.loads(value)
                            except:
                                print ("FAILED to load special content")
                                logger.debug("readPYCDF: Failed to load Object - constructed before v0.2.000?")
                        stream.header[att] = func

        logger.info('readPYCDF: %s Format: %s ' % (filename, cdfformat))

        #if debug:
        #print ("Step2", stream.header)

        variables = cdfdat.cdf_info().get('zVariables')
        timelength = 0

        for key in variables:
            if key.find('time')>=0 or key == 'Epoch':
                # Time column identified
                if not key == 'sectime':
                    ind = KEYLIST.index('time')
                else:
                    ind = KEYLIST.index('sectime')
                try:
                    ttdesc = cdfdat.varinq(key).get('Data_Type_Description')
                    if not ttdesc == 'CDF_TIME_TT2000':
                        print ("WARNING: Time column is not CDF_TIME_TT2000 (found {})".format(ttdesc))
                    col = cdfdat.varget(key)
                    array[ind] = date2num(np.asarray([datetime.utcfromtimestamp(el) for el in cdflib.cdfepoch.unixtime(col)]))
                except:
                    array[ind] = np.asarray([])
            else:
                ind = KEYLIST.index(key)
                addhead = False
                timelength = len(array[0])
                ttdesc = cdfdat.varinq(key).get('Data_Type_Description')
                col = cdfdat.varget(key)

                if not len(col) == timelength and len(col) == 1:
                    array[ind] = np.asarray([col[0]]*timelength)
                    addhead = True
                elif not len(col) == timelength and len(col)>1:
                    array[ind] = np.asarray([])
                else:
                    array[ind] = np.asarray(col)
                    addhead = True
                #print (cdfdat.varattsget(key))
                if addhead:
                    vname = cdfdat.varattsget(key).get('name','')
                    if not vname:
                        vname = cdfdat.varattsget(key).get('FIELDNAM','')
                    vunit = cdfdat.varattsget(key).get('units','')
                    if not vunit:
                        vunit = cdfdat.varattsget(key).get('UNITS','')
                    stream.header['col-'+key.lower()] = vname
                    stream.header['unit-col-'+key.lower()] = vunit

        cdfdat.close()
        del cdfdat

    #print (stream.header)
    return DataStream([LineStruct()], stream.header,np.asarray(array))

def writePYCDF(datastream, filename, **kwargs):
    """
    VARIABLES
        new: use compression variable instead of skipcompression
        compression = 0: skip compression
        compression = 1-9: use this compression factor:
                           9 high compreesion (slow)
                           1 low compression (fast)
               default is 5

    """

    if pyvers and pyvers == 2:
                ch1 = '-'.encode('utf-8') # not working with py3
                ch2 = ''.encode('utf-8')
    else:
                ch1 = '-'
                ch2 = ''

    if not len(datastream.ndarray[0]) > 0 and not len(datastream) > 0:
        return False

    def tt(my_dt_ob):
        ms = my_dt_ob.microsecond/1000.  # fraction
        date_list = [my_dt_ob.year, my_dt_ob.month, my_dt_ob.day, my_dt_ob.hour, my_dt_ob.minute, my_dt_ob.second, ms]
        return date_list

    logger.info("Writing PYCDF Format {}".format(filename))
    mode = kwargs.get('mode')
    skipcompression = kwargs.get('skipcompression')
    compression = kwargs.get('compression')
    version = '1.2'

    if compression == 0: ## temporary solution until all refs to skipcomression are eliminated
        skipcompression = True

    main_cdf_spec = {}
    main_cdf_spec['Compressed'] = False

    leapsecondlastupdate = cdflib.cdfepoch.getLeapSecondLastUpdated()

    if not skipcompression:
        try:
            main_cdf_spec['Compressed'] = True
        except:
            logger.warning("writeIMAGCDF: Compression failed for unknown reason - storing uncompresed data")
            pass

    testname = str(filename+'.cdf')

    if os.path.isfile(testname):
        filename = testname
    if os.path.isfile(filename):
        if mode == 'skip': # skip existing inputs
            exst = read(path_or_url=filename)
            datastream = joinStreams(exst,datastream,extend=True)
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
        elif mode == 'replace': # replace existing inputs
            #print filename
            #### !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            ####  Please note: Replacing requires a lot memory
            #### If memory issues appear then please overwrite existing data
            #### TODO Optimze sorting
            #### !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
            try:
                exst = read(path_or_url=filename)
                datastream = joinStreams(datastream,exst,extend=True)
            except:
                logger.error("writePYCDF: Could not interprete existing data set - aborting")
                sys.exit()
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
        elif mode == 'append':
            #print filename
            exst = read(path_or_url=filename)
            datastream = joinStreams(exst,datastream,extend=True)
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
        else: # overwrite mode
            #print filename
            os.remove(filename)
            mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)
    else:
        mycdf = cdflib.CDF(filename,cdf_spec=main_cdf_spec)

    keylst = datastream._get_key_headers()

    if not 'flag' in keylst:
        keylst.append('flag')
    if not 'comment' in keylst:
        keylst.append('comment')
    if not 'typ' in keylst:
        keylst.append('typ')
    tmpkeylst = ['time']
    tmpkeylst.extend(keylst)
    keylst = tmpkeylst

    headdict = datastream.header
    head, line = [],[]
    globalAttrs = {}

    if not mode == 'append':
        for key in headdict:
            if not key.find('col-') >= 0:
                #print (key, headdict[key])
                if not key in ['DataAbsFunctionObject','DataBaseValues', 'DataFlagList']:
                    globalAttrs[key] = { 0 : str(headdict[key]) }
                else:
                    logger.info("writePYCDF: Found Object in header - pickle and dump ")
                    pfunc = pickle.dumps(headdict[key])
                    globalAttrs[key] = { 0 : pfunc }

    globalAttrs['DataFormat'] = { 0 : 'MagPyCDF{}'.format(version)}

    mycdf.write_globalattrs(globalAttrs)

    def checkEqual3(lst):
        return lst[1:] == lst[:-1]

    for key in keylst:
        var_attrs = {}
        var_spec = {}
        cdfdata = []

        ind = KEYLIST.index(key)
        col = datastream.ndarray[ind]

        if not key in NUMKEYLIST:
            #if not key == 'time':  ### why??
            col = np.asarray(col)

        # Sort out columns only containing nan's
        try:
            test = [elem for elem in col if not isnan(elem)]
            if not len(test) > 0:
                col = np.asarray([])
        except:
            pass

        if not False in checkEqual3(col) and len(col) > 0:
            logger.warning("writePYCDF: Found identical values only for key: %s" % key)
            col = col[:1]

        cdfkey = key.lower()
        if key.find('time') >= 0:
            if cdfkey.endswith('time'):
                if cdfkey == 'time':
                    cdfkey = 'Epoch'
                try: # Might fail for sectime
                    cdfdata = cdflib.cdfepoch.compute_tt2000( [tt(num2date(elem).replace(tzinfo=None)) for elem in col.astype(np.float64)] )
                except:
                    cdfdata = np.asarray([])
                var_spec['Data_Type'] = 33
        elif len(col) > 0:
            if not key in NUMKEYLIST:
                col = list(col)
                col = ['' if el is None else el for el in col]
                col = np.asarray(col) # to get string conversion
                col = list(col) # convert back to list for write_var
                var_spec['Data_Type'] = 51 # CHAR
                var_spec['Num_Elements'] = max([len(i) for i in col])
            else:
                var_spec['Data_Type'] = 45
                col = np.asarray([np.nan if el in [None,ch1] else el for el in col])
                col = col.astype(float)
            cdfdata = col

            var_attrs['name'] = headdict.get('col-'+key,'')       # use 'FIELDNAM' to be conform with NASA / IMAGCDF style ## Version 1.1
            var_attrs['units'] = headdict.get('unit-col-'+key,'')  # use 'UNITS' to be conform with NASA style / IMAGCDF style
            var_attrs['FIELDNAM'] = headdict.get('col-'+key,'')       # use 'FIELDNAM' to be conform with NASA / IMAGCDF style ## Version 1.2
            var_attrs['UNITS'] = headdict.get('unit-col-'+key,'')  # use 'UNITS' to be conform with NASA style / IMAGCDF style ## Version 1.2
            var_attrs['LABLAXIS'] = headdict.get('col-'+key,'') ## Version 1.2

        #if cdfdata is just a single value, then it will be converted into an array
        try:
            test = len(cdfdata)
        except:
            print ("Failed for {}".format(key))
            cdfdata = np.asarray([cdfdata])

        if len(cdfdata) > 0:
            var_spec['Variable'] = cdfkey
            if not var_spec.get('Num_Elements'):
                var_spec['Num_Elements'] = 1
            var_spec['Rec_Vary'] = True # The dimensional sizes, applicable only to rVariables.
            var_spec['Dim_Sizes'] = []

            mycdf.write_var(var_spec, var_attrs=var_attrs, var_data=cdfdata)

    mycdf.close()
    return testname

<|endoftext|>


<commit_before>import os

def _doctest_setup():
    try:
        os.remove("/tmp/open_atomic-example.txt")
    except OSError:
        pass

class open_atomic(object):
    """
    Opens a file for atomic writing by writing to a temporary file, then moving
    the temporary file into place once writing has finished.

    When ``close()`` is called, the temporary file is moved into place,
    overwriting any file which may already exist (except on Windows, see note
    below). If moving the temporary file fails, ``abort()`` will be called *and
    an exception will be raised*.

    If ``abort()`` is called, the temporary file will be removed and the
    ``aborted`` attribute will be set to ``True``. No exception will be raised
    if an error is encountered while removing the temporary file; instead, the
    ``abort_error`` attribute will be set to the exception raised by
    ``os.remove`` (note: on Windows, if ``file.close()`` raises an exception,
    ``abort_error`` will be set to that exception; see implementation of
    ``abort()`` for details).

    By default, ``open_atomic`` will put the temporary file in the same
    directory as the target file:
    ``${dirname(target_file)}/.${basename(target_file)}.temp``. See also the
    ``prefix``, ``suffix``, and ``dir`` arguments to ``open_atomic()``. When
    changing these options, though, remember:

        * Moving files across filesystems is slow, so the temporary file should
          be stored on the same filesystem as the target file.
        * Using a random temporary name is likely a poor idea, as it's more
          likely that temporary files will be left lying around if a process is
          killed and re-started.
        * The temporary file will be blindly overwritten.

    The ``temp_name`` and ``target_name`` attributes store the temporary
    and target file names, and the ``name`` attribute stores the "current"
    name: if the file is still being written, it will store the ``temp_name``,
    and if the temporary file has been moved into place, it will store the
    ``target_name``.

    .. note::

        ``open_atomic`` will appear to work on Windows, but will fail when the
        destination file exists, since ``os.rename`` will not overwrite the
        destination file (specifically, the call to ``close()`` will result in
        an exception being raised, and ``abort()`` being called).

    Example::

        >>> _doctest_setup()
        >>> f = open_atomic("/tmp/open_atomic-example.txt")
        >>> f.temp_name
        '/tmp/.open_atomic-example.txt.temp'
        >>> f.write("Hello, world!")
        >>> (os.path.exists(f.target_name), os.path.exists(f.temp_name))
        (False, True)
        >>> f.close()
        >>> os.path.exists("/tmp/open_atomic-example.txt")
        True

    By default, ``open_atomic`` uses the ``open`` builtin, but this behaviour
    can be changed using the ``opener`` argument::

        >>> import io
        >>> f = open_atomic("/tmp/open_atomic-example.txt",
        ...                opener=io.open,
        ...                mode="w+",
        ...                encoding="utf-8")
        >>> f.write(u"\u1234")
        1L
        >>> f.seek(0)
        0
        >>> f.read()
        u'\u1234'
        >>> f.close()

    """

    def __init__(self, name, prefix=".", suffix=".temp", dir=None, mode="w",
                 opener=open, **open_args):
        self.target_name = name
        self.temp_name = self._get_temp_name(name, prefix, suffix, dir)
        self.fd = opener(self.temp_name, mode, **open_args)
        self.name = self.temp_name
        self.closed = False
        self.aborted = False
        self.abort_error = None

    def _get_temp_name(self, target, prefix, suffix, dir):
        if dir is None:
            dir = os.path.dirname(target)
        return os.path.join(dir, "%s%s%s" %(
            prefix, os.path.basename(target), suffix,
        ))

    def close(self):
        if self.closed:
            return
        try:
            self.fd.close()
            os.rename(self.temp_name, self.target_name)
            self.name = self.target_name
        except:
            try:
                self.abort()
            except:
                pass
            raise
        self.closed = True

    def abort(self):
        try:
            if os.name == "nt":
                # Note: Windows can't remove an open file, so sacrifice some
                # safety and close it before deleting it here. This is only a
                # problem if ``.close()`` raises an exception, which it really
                # shouldn't... But it's probably a better idea to be safe.
                self.fd.close()
            os.remove(self.temp_name)
        except OSError as e:
            self.abort_error = e
        self.fd.close()
        self.closed = True
        self.aborted = True

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        if exc_info[0] is None:
            self.close()
        else:
            self.abort()

    def __getattr__(self, attr):
        return getattr(self.fd, attr)
<commit_msg>Use 'os.replace' when available, use 'self.file' instead of 'self.fd'<commit_after>import os

try:
    replace_func = os.replace
except AttributeError:
    replace_func = os.rename

def _doctest_setup():
    try:
        os.remove("/tmp/open_atomic-example.txt")
    except OSError:
        pass

class open_atomic(object):
    """
    Opens a file for atomic writing by writing to a temporary file, then moving
    the temporary file into place once writing has finished.

    When ``close()`` is called, the temporary file is moved into place,
    overwriting any file which may already exist (except on Windows, see note
    below). If moving the temporary file fails, ``abort()`` will be called *and
    an exception will be raised*.

    If ``abort()`` is called the temporary file will be removed and the
    ``aborted`` attribute will be set to ``True``. No exception will be raised
    if an error is encountered while removing the temporary file; instead, the
    ``abort_error`` attribute will be set to the exception raised by
    ``os.remove`` (note: on Windows, if ``file.close()`` raises an exception,
    ``abort_error`` will be set to that exception; see implementation of
    ``abort()`` for details).

    By default, ``open_atomic`` will put the temporary file in the same
    directory as the target file:
    ``${dirname(target_file)}/.${basename(target_file)}.temp``. See also the
    ``prefix``, ``suffix``, and ``dir`` arguments to ``open_atomic()``. When
    changing these options, remember:

        * The source and the destination must be on the same filesystem,
          otherwise the call to ``os.replace()``/``os.rename()`` may fail (and
          it *will* be much slower than necessary).
        * Using a random temporary name is likely a poor idea, as random names
          will mean it's more likely that temporary files will be left
          abandoned if a process is killed and re-started.
        * The temporary file will be blindly overwritten.

    The ``temp_name`` and ``target_name`` attributes store the temporary
    and target file names, and the ``name`` attribute stores the "current"
    name: if the file is still being written it will store the ``temp_name``,
    and if the temporary file has been moved into place it will store the
    ``target_name``.

    .. note::

        ``open_atomic`` will not work correctly on Windows with Python 2.X or
        Python <= 3.2: the call to ``open_atomic.close()`` will fail when the
        destination file exists (since ``os.rename`` will not overwrite the
        destination file; an exception will be raised and ``abort()`` will be
        called). On Python 3.3 and up ``os.replace`` will be used, which
        will be safe and atomic on both Windows and Unix.

    Example::

        >>> _doctest_setup()
        >>> f = open_atomic("/tmp/open_atomic-example.txt")
        >>> f.temp_name
        '/tmp/.open_atomic-example.txt.temp'
        >>> f.write("Hello, world!")
        >>> (os.path.exists(f.target_name), os.path.exists(f.temp_name))
        (False, True)
        >>> f.close()
        >>> os.path.exists("/tmp/open_atomic-example.txt")
        True

    By default, ``open_atomic`` uses the ``open`` builtin, but this behaviour
    can be changed using the ``opener`` argument::

        >>> import io
        >>> f = open_atomic("/tmp/open_atomic-example.txt",
        ...                opener=io.open,
        ...                mode="w+",
        ...                encoding="utf-8")
        >>> f.write(u"\u1234")
        1L
        >>> f.seek(0)
        0
        >>> f.read()
        u'\u1234'
        >>> f.close()

    """

    def __init__(self, name, prefix=".", suffix=".temp", dir=None, mode="w",
                 opener=open, **open_args):
        self.target_name = name
        self.temp_name = self._get_temp_name(name, prefix, suffix, dir)
        self.file = opener(self.temp_name, mode, **open_args)
        self.name = self.temp_name
        self.closed = False
        self.aborted = False
        self.abort_error = None

    def _get_temp_name(self, target, prefix, suffix, dir):
        if dir is None:
            dir = os.path.dirname(target)
        return os.path.join(dir, "%s%s%s" %(
            prefix, os.path.basename(target), suffix,
        ))

    def close(self):
        if self.closed:
            return
        try:
            self.file.close()
            replace_func(self.temp_name, self.target_name)
            self.name = self.target_name
        except:
            try:
                self.abort()
            except:
                pass
            raise
        self.closed = True

    def abort(self):
        try:
            if os.name == "nt":
                # Note: Windows can't remove an open file, so sacrifice some
                # safety and close it before deleting it here. This is only a
                # problem if ``.close()`` raises an exception, which it really
                # shouldn't... But it's probably a better idea to be safe.
                self.file.close()
            os.remove(self.temp_name)
        except OSError as e:
            self.abort_error = e
        self.file.close()
        self.closed = True
        self.aborted = True

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        if exc_info[0] is None:
            self.close()
        else:
            self.abort()

    def __getattr__(self, attr):
        return getattr(self.file, attr)
<|endoftext|>