<commit_before>"""A tool to create an externally-hosted APK definition JSON from an APK.

For more information see README.md.

"""

import argparse
import base64
from distutils import spawn
import hashlib
import json
import logging
import os.path
import re
import subprocess
import sys
import tempfile
import zipfile

# Enable basic logging to stderr
# logging.basicConfig(level=logging.DEBUG)
logging.basicConfig(level=logging.WARNING)
_log = logging.getLogger(__name__)

class MissingPrerequisiteError(Exception):
  """An exception throw to indicate that a pre-requisit check has failed.

  See `AbstractApkParser` `check_prerequisites`.
  """

class AbstractApkParser(object):
  """Provides an abstract parser that parses an APK file.

  Takes an APK file and returns a dictionary of attributes describing that APK.
  Sub-classes should override the `parse` method.

  Attributes:
    apk_file (str): The path to the APK file being parsed.

  """

  def __init__(self, apk_file):
    """Create an APK parser with a given APK file.

    Args:
      apk_file: The path to the APK file to parse.

    """
    self.apk_file = apk_file
    return

  def check_prerequisites(self):
    """Validates all pre-requisites of the parser.

    For example, a tool expected in the system PATH that is not available.
    This does not guarantee that the parser will not fail, but provides an
    early indication of definite problems.

    Raises:
      MissingPrerequisiteError: If any problems are encountered.

    """
    if not os.path.exists(self.apk_file):
      raise MissingPrerequisiteError("APK file does not exist - "
                                     + str(self.apk_file))
    return

  def parse(self):
    """Parse the APK.  Abstract method that must be overridden.

    Returns:
      A dictionary of elements this parser has extracted from the APK, and
      their values.

    """
    raise NotImplementedError("Abstract method not overriden")

class AaptParser(AbstractApkParser):
  """Parses the APK using the AOSP "aapt" tool.

  Specifically, runs the "aapt dump badging" command, and performs basic
  parsing of the output to extract fields that we need.
  """
  PACKAGE_MATCH_REGEX = re.compile(
      r"\s*package:\s*name='(.*)'\s*"
      r"versionCode='(\d+)'\s*versionName='(.+)'\s*")
  APPLICATION_REGEX = re.compile(
      r"\s*application:\s*label='(.*)'\s*icon='(.*)'\s*")
  APPLICATION_LABEL = re.compile(r"\s*application-label:\s*'(.*)'\s*")
  SDK_VERSION_REGEX = re.compile(r"\s*sdkVersion:\s*'(.*)'\s*")
  MAX_SDK_VERSION_REGEX = re.compile(r"\s*maxSdkVersion:\s*'(.*)'\s*")
  USES_FEATURE_REGEX = re.compile(r"\s*uses-feature:\s+name='(.*)'\s*")
  # Old uses-permission format:
  #     uses-permission:'android.permission.VIBRATE'
  # New uses-permission format:
  #     uses-permission: name='android.permission.VIBRATE'
  #     uses-permission: name='android.permission.WRITE_EXTERNAL_STORAGE'
  #         maxSdkVersion='18'
  USES_PERMISSION_REGEX_OLD = re.compile(r"\s*uses-permission:\s*'(.*)'\s*")
  USES_PERMISSION_REGEX_NEW = re.compile(
      r"\s*uses-permission:\s*name='(.*?)'\s*(?:maxSdkVersion='(.*)'\s*|)")

  def __init__(self, apk_file):
    super(AaptParser, self).__init__(apk_file)
    self.aapt_exe = self.locate_aapt_exe()

  def check_prerequisites(self):
    super(AaptParser, self).check_prerequisites()
    if self.aapt_exe is None:
      raise MissingPrerequisiteError("Couldn't find the aapt binary on "
                                     "system\'s PATH.  This binary is part of "
                                     "the Android developer\'s SDK.  Please "
                                     "ensure it is available on the PATH.")

  def locate_aapt_exe(self):
    return spawn.find_executable("aapt")

  def run_aapt(self):
    return subprocess.check_output(
        "\"" + self.aapt_exe + "\" dump --values badging " + self.apk_file,
        stderr=subprocess.STDOUT,
        shell=True).decode('utf-8').splitlines()

  def parse(self):
    output = {}
    for line in self.run_aapt():
      matches = self.PACKAGE_MATCH_REGEX.match(line)
      if matches:
        _log.info("Matched package")
        output["package_name"] = matches.group(1)
        output["version_code"] = matches.group(2)
        output["version_name"] = matches.group(3)
      matches = self.SDK_VERSION_REGEX.match(line)
      if matches:
        output["minimum_sdk"] = matches.group(1)
      matches = self.MAX_SDK_VERSION_REGEX.match(line)
      if matches:
        output["maximum_sdk"] = matches.group(1)
      matches = self.APPLICATION_LABEL.match(line)
      if matches:
        output["application_label"] = matches.group(1)
      matches = self.APPLICATION_REGEX.match(line)
      if matches:
        # In the case that the explicit "application-label" field is not found
        # in the aapt output, we grab it from the "application" field.
        # (More recent versions of aapt only provide localized versions of
        # application-label in the form "application-label-xx[-XX]".)
        if "application_label" not in output:
          output["application_label"] = matches.group(1)
        output["icon_filename"] = matches.group(2)
      matches = self.USES_FEATURE_REGEX.match(line)
      if matches:
        output.setdefault("uses_feature", []).append(matches.group(1))
      matches = self.USES_PERMISSION_REGEX_OLD.match(line)
      if matches:
        output.setdefault("uses_permission", []).append({"name": matches.group(1)})
      matches = self.USES_PERMISSION_REGEX_NEW.match(line)
      if matches:
        new_permission = {"name": matches.group(1)}
        try:
          if matches.group(2) is not None:
            new_permission.update({"maxSdkVersion": matches.group(2)})
        except IndexError:
          # No maxSdkVersion - that's OK, it's not mandatory
          pass
        output.setdefault("uses_permission", []).append(new_permission)
    return output

class FileParser(AbstractApkParser):
  """Parses properties of the APK file as a system file.

  """

  def parse(self):
    output = {}
    output["file_size"] = os.path.getsize(self.apk_file)

    with open(self.apk_file, "rb") as f:
      output["file_sha1_base64"] = base64.b64encode(
          hashlib.sha1(f.read()).digest()).decode('ascii')

    with open(self.apk_file, "rb") as f:
      output["file_sha256_base64"] = base64.b64encode(
          hashlib.sha256(f.read()).digest()).decode('ascii')

    return output

class JarFileValidationParser(AbstractApkParser):
  """Parser that validates that the APK is a valid zip-aligned JAR.

  Note:
    This parser doesn't actually return any information from the APK,
    (it returns an empty dictionary), but it produces errors if the
    APK fails validation.

  """

  def __init__(self, apk_file):
    super(JarFileValidationParser, self).__init__(apk_file)
    self.zipalign_exe = self.locate_zipalign_exe()
    self.jarsigner_exe = self.locate_jarsigner_exe()

  def locate_zipalign_exe(self):
    return spawn.find_executable("zipalign")

  def locate_jarsigner_exe(self):
    return spawn.find_executable("jarsigner")

  def check_prerequisites(self):
    super(JarFileValidationParser, self).check_prerequisites()
    if self.zipalign_exe is None:
      raise MissingPrerequisiteError("Couldn't find zipalign binary in "
                                     "system's PATH.  This binary is needed "
                                     "to validate the APK.")
    if self.jarsigner_exe is None:
      raise MissingPrerequisiteError("Couldn't find jarsigner binary in "
                                     "system's PATH.  This binary is needed "
                                     "to validate the APK.")
    return

  def parse(self):
    # Validate that the zip is correctly aligned.
    try:
      subprocess.check_call("\"" + self.zipalign_exe + "\" -c 4 " + self.apk_file,
                            shell=True)
    except subprocess.CalledProcessError as e:
      raise Exception("Error: Zip alignment is incorrect", e)

    # Validate that the jar is signed correctly.
    with open(os.devnull, "w") as dev_null:
      try:
        subprocess.check_call("\"" + self.jarsigner_exe + "\" -verify " + self.apk_file,
                              stdout=dev_null,
                              shell=True)
      except subprocess.CalledProcessError as e:
        raise Exception("Error: JAR signature doesn't validate correctly", e)

    # No new data parsed from APK, return an empty dictionary.
    return {}

class IconParser(AbstractApkParser):
  """Parses the icon from the file as base64 encoded binary.

  Attributes:
    icon_filename (str) : Filename of the icon within the APK.

  """

  def __init__(self, apk_file, icon_filename):
    super(IconParser, self).__init__(apk_file)
    self.icon_filename = icon_filename

  def check_prerequisites(self):
    super(IconParser, self).check_prerequisites()
    if self.icon_filename is None:
      raise MissingPrerequisiteError("Couldn't find icon in APK")

  def parse(self):
    output = {}
    jar_zip = zipfile.ZipFile(self.apk_file)
    icon_file_bytes = jar_zip.read(self.icon_filename)
    output["icon_base64"] = base64.b64encode(icon_file_bytes).decode('ascii')
    return output

class CertificateParser(AbstractApkParser):
  """Parses the signing certificate chain from the APK.

  """

  def __init__(self, apk_file):
    super(CertificateParser, self).__init__(apk_file)
    self.openssl_exe = self.locate_openssl_exe()

  def locate_openssl_exe(self):
    return spawn.find_executable("openssl")

  def openssl_convert_rsa_cert(self, infile_path, outfile_path):
    subprocess.check_call("\"" + self.openssl_exe + "\" pkcs7 -in "
                          + infile_path
                          + " -print_certs -inform DER -out "
                          + outfile_path, shell=True)

  def check_prerequisites(self):
    # We use a command line tool rather than the openssl library to
    # simplify install for developers not experienced with Python.
    super(CertificateParser, self).check_prerequisites()
    if self.openssl_exe is None:
      raise MissingPrerequisiteError("Couldn't find openssl commandline tool "
                                     "in system PATH.")

  def parse(self):
    output = {}
    jar_zip = zipfile.ZipFile(self.apk_file)

    # Find all the files in the JAR manifest (which includes the certificate)
    # and extract into a temporary directory.
    temp_dir = tempfile.mkdtemp()
    for ii in jar_zip.namelist():
      _log.info(ii)
    _log.info("Found " + str(len(jar_zip.namelist())) + " files in the APK")
    manifest_files = [ii for ii in jar_zip.namelist()
                      if ii.startswith("META-INF")]
    for f in manifest_files:
      _log.info("Found file in manfiest folder: " + f)
    jar_zip.extractall(path=temp_dir, members=manifest_files)

    # Look at each file to try and find the RSA certificate file (we expect
    # only one, and throw an error if there are multiple).
    with tempfile.NamedTemporaryFile(
        mode="r+b", delete=False) as temp_out_file:
      _log.info("Writing to temporary output file" + temp_out_file.name)
      for path, _, files in os.walk(temp_dir):
        for f in files:
          if not f.endswith(".RSA"):
            continue
          if output.get("certificate_base64") is not None:
            raise Exception("Multiple RSA keys - APK should only be signed "
                            "by a single jarsigner.")

          # Found the RSA certificate.  Use the openssl commandline tool to
          # extract details.
          _log.info("Found RSA file: " + os.path.join(path, f))
          self.openssl_convert_rsa_cert(os.path.join(path, f),
                                        temp_out_file.name)

          # Mark the output to indicate we have found a certificate, so that
          # even if we fail to parse it properly it is counted.
          output["certificate_base64"] = []

          # Certificates have been dumped, in order (if a cert chain existed)
          # in base64 format, which is what we need.
          certificate = None
          for ii in temp_out_file:
            ii = ii.decode('ascii').strip()
            if re.match(r".*-+BEGIN\s+CERTIFICATE-+.*", ii):
              _log.debug("Begin certificate line")
              certificate = ""
            elif re.match(r".*-+END\s+CERTIFICATE-+.*", ii):
              _log.debug("End certificate line")
              output["certificate_base64"].append(certificate)
              certificate = None
            elif certificate is not None:
              certificate += ii
            else:
              _log.debug("Skipping non-cert line: " + ii)
    return output

class DataExtractor(object):
  """Class to extract the required information from the apk and hosting url.
  """

  def __init__(self, apk, externallyhostedurl):
    self.apk = apk
    self.externallyhostedurl = externallyhostedurl

  def parse(self):
    """Parse the given apk and hosting url.

    This extracts the required information and returns a dict which needs to
    be formatted as JSON to form the file that can be uploaded in the developer
    console."""

    if not os.path.exists(self.apk):
      raise MissingPrerequisiteError("Could not find APK " + self.apk)

    # Create a list of parser that provides all the data we need for our
    # externally-hosted APK JSON.
    parsers = [AaptParser(self.apk), FileParser(self.apk),
               CertificateParser(self.apk), JarFileValidationParser(self.apk)]

    # Validate the system setup
    for parser in parsers:
      parser.check_prerequisites()

    # Parse the APK
    apk_properties = {}
    for parser in parsers:
      apk_properties.update(parser.parse())

    # Also add the icon (this relies on data parsed in the previous parser).
    icon_parser = IconParser(self.apk, apk_properties.get("icon_filename"))
    icon_parser.check_prerequisites()
    apk_properties.update(icon_parser.parse())

    # Add in the externally-hosted URL and print to stdout.
    apk_properties["externally_hosted_url"] = self.externallyhostedurl
    return apk_properties

def main():
  """Print an externally-hosted APK JSON definition to stdout.

  """
  # Define flags and help text.
  arg_parser = argparse.ArgumentParser(
      description="For a given APK, create an externally hosted APK "
      "definition file.")
  arg_parser.add_argument("--apk", dest="apk_file_name", required=True)
  arg_parser.add_argument("--externallyHostedUrl",
                          dest="externally_hosted_url", required=True)
  args = arg_parser.parse_args()

  extractor = DataExtractor(args.apk_file_name, args.externally_hosted_url)
  print(json.dumps(extractor.parse(), indent=1))
  return

if __name__ == "__main__":
  main()
<commit_msg>Clamping min_sdk version to 26 since otherwise the upload will fail.<commit_after>"""A tool to create an externally-hosted APK definition JSON from an APK.

For more information see README.md.

"""

import argparse
import base64
from distutils import spawn
import hashlib
import json
import logging
import os.path
import re
import subprocess
import sys
import tempfile
import zipfile

# Google Play Console now has a minimum version of 26.
MIN_ALLOWED_SDK_VERSION = 26

# Enable basic logging to stderr
# logging.basicConfig(level=logging.DEBUG)
logging.basicConfig(level=logging.WARNING)
_log = logging.getLogger(__name__)

class MissingPrerequisiteError(Exception):
  """An exception throw to indicate that a pre-requisit check has failed.

  See `AbstractApkParser` `check_prerequisites`.
  """

class AbstractApkParser(object):
  """Provides an abstract parser that parses an APK file.

  Takes an APK file and returns a dictionary of attributes describing that APK.
  Sub-classes should override the `parse` method.

  Attributes:
    apk_file (str): The path to the APK file being parsed.

  """

  def __init__(self, apk_file):
    """Create an APK parser with a given APK file.

    Args:
      apk_file: The path to the APK file to parse.

    """
    self.apk_file = apk_file
    return

  def check_prerequisites(self):
    """Validates all pre-requisites of the parser.

    For example, a tool expected in the system PATH that is not available.
    This does not guarantee that the parser will not fail, but provides an
    early indication of definite problems.

    Raises:
      MissingPrerequisiteError: If any problems are encountered.

    """
    if not os.path.exists(self.apk_file):
      raise MissingPrerequisiteError("APK file does not exist - "
                                     + str(self.apk_file))
    return

  def parse(self):
    """Parse the APK.  Abstract method that must be overridden.

    Returns:
      A dictionary of elements this parser has extracted from the APK, and
      their values.

    """
    raise NotImplementedError("Abstract method not overriden")

class AaptParser(AbstractApkParser):
  """Parses the APK using the AOSP "aapt" tool.

  Specifically, runs the "aapt dump badging" command, and performs basic
  parsing of the output to extract fields that we need.
  """
  PACKAGE_MATCH_REGEX = re.compile(
      r"\s*package:\s*name='(.*)'\s*"
      r"versionCode='(\d+)'\s*versionName='(.+)'\s*")
  APPLICATION_REGEX = re.compile(
      r"\s*application:\s*label='(.*)'\s*icon='(.*)'\s*")
  APPLICATION_LABEL = re.compile(r"\s*application-label:\s*'(.*)'\s*")
  SDK_VERSION_REGEX = re.compile(r"\s*sdkVersion:\s*'(.*)'\s*")
  MAX_SDK_VERSION_REGEX = re.compile(r"\s*maxSdkVersion:\s*'(.*)'\s*")
  USES_FEATURE_REGEX = re.compile(r"\s*uses-feature:\s+name='(.*)'\s*")
  # Old uses-permission format:
  #     uses-permission:'android.permission.VIBRATE'
  # New uses-permission format:
  #     uses-permission: name='android.permission.VIBRATE'
  #     uses-permission: name='android.permission.WRITE_EXTERNAL_STORAGE'
  #         maxSdkVersion='18'
  USES_PERMISSION_REGEX_OLD = re.compile(r"\s*uses-permission:\s*'(.*)'\s*")
  USES_PERMISSION_REGEX_NEW = re.compile(
      r"\s*uses-permission:\s*name='(.*?)'\s*(?:maxSdkVersion='(.*)'\s*|)")

  def __init__(self, apk_file):
    super(AaptParser, self).__init__(apk_file)
    self.aapt_exe = self.locate_aapt_exe()

  def check_prerequisites(self):
    super(AaptParser, self).check_prerequisites()
    if self.aapt_exe is None:
      raise MissingPrerequisiteError("Couldn't find the aapt binary on "
                                     "system\'s PATH.  This binary is part of "
                                     "the Android developer\'s SDK.  Please "
                                     "ensure it is available on the PATH.")

  def locate_aapt_exe(self):
    return spawn.find_executable("aapt")

  def run_aapt(self):
    return subprocess.check_output(
        "\"" + self.aapt_exe + "\" dump --values badging " + self.apk_file,
        stderr=subprocess.STDOUT,
        shell=True).decode('utf-8').splitlines()

  def parse(self):
    output = {}
    for line in self.run_aapt():
      matches = self.PACKAGE_MATCH_REGEX.match(line)
      if matches:
        _log.info("Matched package")
        output["package_name"] = matches.group(1)
        output["version_code"] = matches.group(2)
        output["version_name"] = matches.group(3)
      matches = self.SDK_VERSION_REGEX.match(line)
      if matches:
        min_sdk = int(matches.group(1), 10)
        if min_sdk < MIN_ALLOWED_SDK_VERSION:
          _log.warn("Bumping min sdk from %s to %s",
                    min_sdk, MIN_ALLOWED_SDK_VERSION)
          min_sdk = MIN_ALLOWED_SDK_VERSION
        output["minimum_sdk"] = str(min_sdk)
      matches = self.MAX_SDK_VERSION_REGEX.match(line)
      if matches:
        output["maximum_sdk"] = matches.group(1)
      matches = self.APPLICATION_LABEL.match(line)
      if matches:
        output["application_label"] = matches.group(1)
      matches = self.APPLICATION_REGEX.match(line)
      if matches:
        # In the case that the explicit "application-label" field is not found
        # in the aapt output, we grab it from the "application" field.
        # (More recent versions of aapt only provide localized versions of
        # application-label in the form "application-label-xx[-XX]".)
        if "application_label" not in output:
          output["application_label"] = matches.group(1)
        output["icon_filename"] = matches.group(2)
      matches = self.USES_FEATURE_REGEX.match(line)
      if matches:
        output.setdefault("uses_feature", []).append(matches.group(1))
      matches = self.USES_PERMISSION_REGEX_OLD.match(line)
      if matches:
        output.setdefault("uses_permission", []).append({"name": matches.group(1)})
      matches = self.USES_PERMISSION_REGEX_NEW.match(line)
      if matches:
        new_permission = {"name": matches.group(1)}
        try:
          if matches.group(2) is not None:
            new_permission.update({"maxSdkVersion": matches.group(2)})
        except IndexError:
          # No maxSdkVersion - that's OK, it's not mandatory
          pass
        output.setdefault("uses_permission", []).append(new_permission)
    return output

class FileParser(AbstractApkParser):
  """Parses properties of the APK file as a system file.

  """

  def parse(self):
    output = {}
    output["file_size"] = os.path.getsize(self.apk_file)

    with open(self.apk_file, "rb") as f:
      output["file_sha1_base64"] = base64.b64encode(
          hashlib.sha1(f.read()).digest()).decode('ascii')

    with open(self.apk_file, "rb") as f:
      output["file_sha256_base64"] = base64.b64encode(
          hashlib.sha256(f.read()).digest()).decode('ascii')

    return output

class JarFileValidationParser(AbstractApkParser):
  """Parser that validates that the APK is a valid zip-aligned JAR.

  Note:
    This parser doesn't actually return any information from the APK,
    (it returns an empty dictionary), but it produces errors if the
    APK fails validation.

  """

  def __init__(self, apk_file):
    super(JarFileValidationParser, self).__init__(apk_file)
    self.zipalign_exe = self.locate_zipalign_exe()
    self.jarsigner_exe = self.locate_jarsigner_exe()

  def locate_zipalign_exe(self):
    return spawn.find_executable("zipalign")

  def locate_jarsigner_exe(self):
    return spawn.find_executable("jarsigner")

  def check_prerequisites(self):
    super(JarFileValidationParser, self).check_prerequisites()
    if self.zipalign_exe is None:
      raise MissingPrerequisiteError("Couldn't find zipalign binary in "
                                     "system's PATH.  This binary is needed "
                                     "to validate the APK.")
    if self.jarsigner_exe is None:
      raise MissingPrerequisiteError("Couldn't find jarsigner binary in "
                                     "system's PATH.  This binary is needed "
                                     "to validate the APK.")
    return

  def parse(self):
    # Validate that the zip is correctly aligned.
    try:
      subprocess.check_call("\"" + self.zipalign_exe + "\" -c 4 " + self.apk_file,
                            shell=True)
    except subprocess.CalledProcessError as e:
      raise Exception("Error: Zip alignment is incorrect", e)

    # Validate that the jar is signed correctly.
    with open(os.devnull, "w") as dev_null:
      try:
        subprocess.check_call("\"" + self.jarsigner_exe + "\" -verify " + self.apk_file,
                              stdout=dev_null,
                              shell=True)
      except subprocess.CalledProcessError as e:
        raise Exception("Error: JAR signature doesn't validate correctly", e)

    # No new data parsed from APK, return an empty dictionary.
    return {}

class IconParser(AbstractApkParser):
  """Parses the icon from the file as base64 encoded binary.

  Attributes:
    icon_filename (str) : Filename of the icon within the APK.

  """

  def __init__(self, apk_file, icon_filename):
    super(IconParser, self).__init__(apk_file)
    self.icon_filename = icon_filename

  def check_prerequisites(self):
    super(IconParser, self).check_prerequisites()
    if self.icon_filename is None:
      raise MissingPrerequisiteError("Couldn't find icon in APK")

  def parse(self):
    output = {}
    jar_zip = zipfile.ZipFile(self.apk_file)
    icon_file_bytes = jar_zip.read(self.icon_filename)
    output["icon_base64"] = base64.b64encode(icon_file_bytes).decode('ascii')
    return output

class CertificateParser(AbstractApkParser):
  """Parses the signing certificate chain from the APK.

  """

  def __init__(self, apk_file):
    super(CertificateParser, self).__init__(apk_file)
    self.openssl_exe = self.locate_openssl_exe()

  def locate_openssl_exe(self):
    return spawn.find_executable("openssl")

  def openssl_convert_rsa_cert(self, infile_path, outfile_path):
    subprocess.check_call("\"" + self.openssl_exe + "\" pkcs7 -in "
                          + infile_path
                          + " -print_certs -inform DER -out "
                          + outfile_path, shell=True)

  def check_prerequisites(self):
    # We use a command line tool rather than the openssl library to
    # simplify install for developers not experienced with Python.
    super(CertificateParser, self).check_prerequisites()
    if self.openssl_exe is None:
      raise MissingPrerequisiteError("Couldn't find openssl commandline tool "
                                     "in system PATH.")

  def parse(self):
    output = {}
    jar_zip = zipfile.ZipFile(self.apk_file)

    # Find all the files in the JAR manifest (which includes the certificate)
    # and extract into a temporary directory.
    temp_dir = tempfile.mkdtemp()
    for ii in jar_zip.namelist():
      _log.info(ii)
    _log.info("Found " + str(len(jar_zip.namelist())) + " files in the APK")
    manifest_files = [ii for ii in jar_zip.namelist()
                      if ii.startswith("META-INF")]
    for f in manifest_files:
      _log.info("Found file in manfiest folder: " + f)
    jar_zip.extractall(path=temp_dir, members=manifest_files)

    # Look at each file to try and find the RSA certificate file (we expect
    # only one, and throw an error if there are multiple).
    with tempfile.NamedTemporaryFile(
        mode="r+b", delete=False) as temp_out_file:
      _log.info("Writing to temporary output file" + temp_out_file.name)
      for path, _, files in os.walk(temp_dir):
        for f in files:
          if not f.endswith(".RSA"):
            continue
          if output.get("certificate_base64") is not None:
            raise Exception("Multiple RSA keys - APK should only be signed "
                            "by a single jarsigner.")

          # Found the RSA certificate.  Use the openssl commandline tool to
          # extract details.
          _log.info("Found RSA file: " + os.path.join(path, f))
          self.openssl_convert_rsa_cert(os.path.join(path, f),
                                        temp_out_file.name)

          # Mark the output to indicate we have found a certificate, so that
          # even if we fail to parse it properly it is counted.
          output["certificate_base64"] = []

          # Certificates have been dumped, in order (if a cert chain existed)
          # in base64 format, which is what we need.
          certificate = None
          for ii in temp_out_file:
            ii = ii.decode('ascii').strip()
            if re.match(r".*-+BEGIN\s+CERTIFICATE-+.*", ii):
              _log.debug("Begin certificate line")
              certificate = ""
            elif re.match(r".*-+END\s+CERTIFICATE-+.*", ii):
              _log.debug("End certificate line")
              output["certificate_base64"].append(certificate)
              certificate = None
            elif certificate is not None:
              certificate += ii
            else:
              _log.debug("Skipping non-cert line: " + ii)
    return output

class DataExtractor(object):
  """Class to extract the required information from the apk and hosting url.
  """

  def __init__(self, apk, externallyhostedurl):
    self.apk = apk
    self.externallyhostedurl = externallyhostedurl

  def parse(self):
    """Parse the given apk and hosting url.

    This extracts the required information and returns a dict which needs to
    be formatted as JSON to form the file that can be uploaded in the developer
    console."""

    if not os.path.exists(self.apk):
      raise MissingPrerequisiteError("Could not find APK " + self.apk)

    # Create a list of parser that provides all the data we need for our
    # externally-hosted APK JSON.
    parsers = [AaptParser(self.apk), FileParser(self.apk),
               CertificateParser(self.apk), JarFileValidationParser(self.apk)]

    # Validate the system setup
    for parser in parsers:
      parser.check_prerequisites()

    # Parse the APK
    apk_properties = {}
    for parser in parsers:
      apk_properties.update(parser.parse())

    # Also add the icon (this relies on data parsed in the previous parser).
    icon_parser = IconParser(self.apk, apk_properties.get("icon_filename"))
    icon_parser.check_prerequisites()
    apk_properties.update(icon_parser.parse())

    # Add in the externally-hosted URL and print to stdout.
    apk_properties["externally_hosted_url"] = self.externallyhostedurl
    return apk_properties

def main():
  """Print an externally-hosted APK JSON definition to stdout.

  """
  # Define flags and help text.
  arg_parser = argparse.ArgumentParser(
      description="For a given APK, create an externally hosted APK "
      "definition file.")
  arg_parser.add_argument("--apk", dest="apk_file_name", required=True)
  arg_parser.add_argument("--externallyHostedUrl",
                          dest="externally_hosted_url", required=True)
  args = arg_parser.parse_args()

  extractor = DataExtractor(args.apk_file_name, args.externally_hosted_url)
  print(json.dumps(extractor.parse(), indent=1))
  return

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


<commit_before>"""Extractor for fanfiction.net."""

from fanfic_scraper.base_fanfic import BaseFanfic, BaseChapter
from urllib.parse import urlparse, urljoin
import requests
import bs4 as bsoup
from collections import defaultdict
import re
import os
import shutil
import textwrap
import datetime
from random import shuffle, uniform
from copy import deepcopy
from time import sleep

def chapter_nav(tag):
    test = (tag.name == 'select')
    test = (test and 'chap_select' in tag['id'])

    return test

class FanfictionNetFanfic(BaseFanfic):

    def get_story_url(self, storyid):
        return 'https://www.fanfiction.net/s/'+storyid

    def extract_chapters(self):
        """Extract chapters function (backbone)."""
        fanfic_name = self.name
        url = self.url
        self.title = url.split('/')[-1]
        urlscheme = urlparse(url)

        #set story_id from url
        self.fanfic_id = urlscheme.path.split('/')[2]

        # Get chapters
        r = requests.get(url, verify=self.verify_https)
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        chapters = defaultdict(FanfictionNetChapter)

        try:

            chapter_list = soup.find_all(chapter_nav)[0]

            for link in chapter_list:
                chapter_num = int(link.get('value'))
                chapter_link = urljoin(
                    urlscheme.scheme + "://" + urlscheme.netloc,
                    "s/" + self.fanfic_id + "/" + str(chapter_num))
                chapters[chapter_num] = FanfictionNetChapter(
                    self, chapter_num, chapter_link)

            return chapters

        except:
            return chapters

    def get_update_date(self):
        r = requests.get(self.url, verify=self.verify_https)
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            span = div.find_all('span')[-2]
            ts = span.get('data-xutime')
            update_date = datetime.datetime.fromtimestamp(float(ts)).strftime('%Y-%m-%d')

            break

        return update_date

class FanfictionNetChapter(BaseChapter):
    """Base chapter class."""

    def get_fanfic_title(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            title = div.find_all('b')[0].get_text()
            break

        return title

    def get_fanfic_author(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            for link in div.find_all('a'):
                if '/u' in link.get('href'):
                    author = link.get_text()
                    break

        return author

    def get_fanfic_category(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'pre_story_links'}):
            for link in div.find_all('a'):
                category = link.get_text()
                break

        return category

    def get_fanfic_description(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            description = div.find_all('div')[1].get_text()
            break

        return description

    def get_update_date(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            span = div.find_all('span')[-2]
            ts = span.get('data-xutime')
            update_date = datetime.datetime.fromtimestamp(float(ts)).strftime('%Y-%m-%d')

            break

        return update_date

    def get_publish_date(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            span = div.find_all('span')[-1]
            ts = span.get('data-xutime')
            publish_date = datetime.datetime.fromtimestamp(float(ts)).strftime('%Y-%m-%d')

            break

        return publish_date

    def get_chapter_title(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        chapter_list = soup.find_all(chapter_nav)[0]
        for link in chapter_list:
            if int(link.get('value')) == self.chapter_num:
                chapter_title = link.get_text()
                break

        return chapter_title

    def get_chapter_count(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        chapter_list = soup.find_all(chapter_nav)[0]

        return len(chapter_list)

    def get_chapter_html(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        story = soup.find_all('div',{'id':'storytext'})[0]

        return str(story)

    def render_p(self,value):
        return '
'+value+'
'

    def story_info(self):
        r = send_request(self.fanfic_url, self.verify_https)

        title = self.get_fanfic_title(r)
        author = self.get_fanfic_author(r)
        category = self.get_fanfic_category(r)
        desc = self.get_fanfic_description(r)
        update_date = self.get_update_date(r)
        publish_date = self.get_publish_date(r)
        chapter_count = self.get_chapter_count(r)

        info = {}

        info['StoryId'] = self.fanfic_id
        info['Title'] = title
        info['Author'] = author
        info['Category'] = category
        info['Description'] = desc
        info['Publish_Date'] = publish_date
        info['Update_Date'] = update_date
        info['Count'] = chapter_count

        return info

    def download_chapter(self):

        filename = self.fanfic_name+'-%03d.htm' % (self.chapter_num)

        r = requests.get(self.chapter_url, verify=self.verify_https)

        title = self.get_fanfic_title(r)
        author = self.get_fanfic_author(r)
        category = self.get_fanfic_category(r)
        desc = self.get_fanfic_description(r)
        chapter_title = self.get_chapter_title(r)
        update_date = self.get_update_date(r)
        publish_date = self.get_publish_date(r)
        chapter_count = self.get_chapter_count(r)
        story = self.get_chapter_html(r)

        #print(title)
        #print(author)
        #print(category)
        #print("Summary: ", textwrap.fill(desc))
        #print('Chapter '+chapter_title)
        #print('Published: '+publish_date)
        #print('Updated: '+update_date)
        #print(chapter_count)
        #print(story)

        target = os.path.join(self.fanfic_download_location,filename)

        if os.path.isfile(target):
            os.remove(target)

        f1 = open(target, "w")

        f1.write('')
        f1.write('')
        f1.write(self.render_p(title))
        f1.write(self.render_p(author))
        f1.write(self.render_p(category))
        f1.write(self.render_p('Summary: '+desc))
        f1.write(self.render_p('Chapter '+chapter_title))
        if self.chapter_num == 1:
            f1.write(self.render_p('Published: '+publish_date))
        if self.chapter_num == chapter_count:
            f1.write(self.render_p('Updated: '+update_date))
        f1.write(self.render_p('========='))
        f1.write(story)

        f1.flush()
        os.fsync(f1.fileno())
        f1.close

<commit_msg>Bug Fix: Correct request FFNet Story Info.<commit_after>"""Extractor for fanfiction.net."""

from fanfic_scraper.base_fanfic import BaseFanfic, BaseChapter
from urllib.parse import urlparse, urljoin
import requests
import bs4 as bsoup
from collections import defaultdict
import re
import os
import shutil
import textwrap
import datetime
from random import shuffle, uniform
from copy import deepcopy
from time import sleep

def chapter_nav(tag):
    test = (tag.name == 'select')
    test = (test and 'chap_select' in tag['id'])

    return test

class FanfictionNetFanfic(BaseFanfic):

    def get_story_url(self, storyid):
        return 'https://www.fanfiction.net/s/'+storyid

    def extract_chapters(self):
        """Extract chapters function (backbone)."""
        fanfic_name = self.name
        url = self.url
        self.title = url.split('/')[-1]
        urlscheme = urlparse(url)

        #set story_id from url
        self.fanfic_id = urlscheme.path.split('/')[2]

        # Get chapters
        r = requests.get(url, verify=self.verify_https)
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        chapters = defaultdict(FanfictionNetChapter)

        try:

            chapter_list = soup.find_all(chapter_nav)[0]

            for link in chapter_list:
                chapter_num = int(link.get('value'))
                chapter_link = urljoin(
                    urlscheme.scheme + "://" + urlscheme.netloc,
                    "s/" + self.fanfic_id + "/" + str(chapter_num))
                chapters[chapter_num] = FanfictionNetChapter(
                    self, chapter_num, chapter_link)

            return chapters

        except:
            return chapters

    def get_update_date(self):
        r = requests.get(self.url, verify=self.verify_https)
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            span = div.find_all('span')[-2]
            ts = span.get('data-xutime')
            update_date = datetime.datetime.fromtimestamp(float(ts)).strftime('%Y-%m-%d')

            break

        return update_date

class FanfictionNetChapter(BaseChapter):
    """Base chapter class."""

    def get_fanfic_title(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            title = div.find_all('b')[0].get_text()
            break

        return title

    def get_fanfic_author(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            for link in div.find_all('a'):
                if '/u' in link.get('href'):
                    author = link.get_text()
                    break

        return author

    def get_fanfic_category(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'pre_story_links'}):
            for link in div.find_all('a'):
                category = link.get_text()
                break

        return category

    def get_fanfic_description(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            description = div.find_all('div')[1].get_text()
            break

        return description

    def get_update_date(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            span = div.find_all('span')[-2]
            ts = span.get('data-xutime')
            update_date = datetime.datetime.fromtimestamp(float(ts)).strftime('%Y-%m-%d')

            break

        return update_date

    def get_publish_date(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        for div in soup.find_all('div', {'id':'profile_top'}):
            span = div.find_all('span')[-1]
            ts = span.get('data-xutime')
            publish_date = datetime.datetime.fromtimestamp(float(ts)).strftime('%Y-%m-%d')

            break

        return publish_date

    def get_chapter_title(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        chapter_list = soup.find_all(chapter_nav)[0]
        for link in chapter_list:
            if int(link.get('value')) == self.chapter_num:
                chapter_title = link.get_text()
                break

        return chapter_title

    def get_chapter_count(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        chapter_list = soup.find_all(chapter_nav)[0]

        return len(chapter_list)

    def get_chapter_html(self,r):
        soup = bsoup.BeautifulSoup(r.text, 'html5lib')

        story = soup.find_all('div',{'id':'storytext'})[0]

        return str(story)

    def render_p(self,value):
        return '
'+value+'
'

    def story_info(self):
        r = requests.get(self.chapter_url, verify=self.verify_https)

        title = self.get_fanfic_title(r)
        author = self.get_fanfic_author(r)
        category = self.get_fanfic_category(r)
        desc = self.get_fanfic_description(r)
        update_date = self.get_update_date(r)
        publish_date = self.get_publish_date(r)
        chapter_count = self.get_chapter_count(r)

        info = {}

        info['StoryId'] = self.fanfic_id
        info['Title'] = title
        info['Author'] = author
        info['Category'] = category
        info['Description'] = desc
        info['Publish_Date'] = publish_date
        info['Update_Date'] = update_date
        info['Count'] = chapter_count

        return info

    def download_chapter(self):

        filename = self.fanfic_name+'-%03d.htm' % (self.chapter_num)

        r = requests.get(self.chapter_url, verify=self.verify_https)

        title = self.get_fanfic_title(r)
        author = self.get_fanfic_author(r)
        category = self.get_fanfic_category(r)
        desc = self.get_fanfic_description(r)
        chapter_title = self.get_chapter_title(r)
        update_date = self.get_update_date(r)
        publish_date = self.get_publish_date(r)
        chapter_count = self.get_chapter_count(r)
        story = self.get_chapter_html(r)

        #print(title)
        #print(author)
        #print(category)
        #print("Summary: ", textwrap.fill(desc))
        #print('Chapter '+chapter_title)
        #print('Published: '+publish_date)
        #print('Updated: '+update_date)
        #print(chapter_count)
        #print(story)

        target = os.path.join(self.fanfic_download_location,filename)

        if os.path.isfile(target):
            os.remove(target)

        f1 = open(target, "w")

        f1.write('')
        f1.write('')
        f1.write(self.render_p(title))
        f1.write(self.render_p(author))
        f1.write(self.render_p(category))
        f1.write(self.render_p('Summary: '+desc))
        f1.write(self.render_p('Chapter '+chapter_title))
        if self.chapter_num == 1:
            f1.write(self.render_p('Published: '+publish_date))
        if self.chapter_num == chapter_count:
            f1.write(self.render_p('Updated: '+update_date))
        f1.write(self.render_p('========='))
        f1.write(story)

        f1.flush()
        os.fsync(f1.fileno())
        f1.close

<|endoftext|>


<commit_before>import numpy as np
from matplotlib import gridspec
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import FixedLocator
from seyiranalyzer.AudioSeyirAnalyzer import AudioSeyirAnalyzer

class Plotter(object):
    @staticmethod
    def plot_audio_features(pitch=None, pitch_distribution=None,
                            sections=None, notes=None, note_models=None,
                            melodic_progression=None):

        # create the figure and all the subplots with the shared axis specified
        fig, ax1, ax2, ax3, ax4 = Plotter._create_figure()

        # plot the pitch track and the performed notes to the first subplot
        Plotter._subplot_pitch_notes(ax1, notes, pitch)

        # plot the pitch distribution and the note models to the second subplot
        Plotter._plot_pitch_dist_note_models(ax2, note_models,
                                             pitch_distribution)

        # plot the melodic progression to the third subplot
        Plotter._plot_melodic_progression(ax3, melodic_progression, pitch,
                                          pitch_distribution)

        # plot the sections to the third subplot, onto the melodic progression
        Plotter._plot_sections(ax1, ax3, ax4, sections)

        return fig, (ax1, ax2, ax3, ax4)

    @staticmethod
    def _create_figure():
        # create the figure with four subplots with different size
        # - 1st is for the predominant melody and performed notes
        # - 2nd is the pitch distribution and note models, it shares the y
        # axis with the 1st
        # - 3rd is the melodic progression, it shares the x axis with the 1st
        # - 4th is for the sections, it is on top the 3rd
        fig = plt.figure()
        gs = gridspec.GridSpec(2, 2, width_ratios=[6, 1], height_ratios=[4, 1])
        ax1 = fig.add_subplot(gs[0])  # pitch and notes
        ax2 = fig.add_subplot(gs[1], sharey=ax1)  # pitch dist. and note models
        ax4 = fig.add_subplot(gs[2])  # sections
        ax3 = plt.twiny(ax4)  # melodic progression
        ax1.get_shared_x_axes().join(ax1, ax3)
        fig.subplots_adjust(hspace=0, wspace=0)
        return fig, ax1, ax2, ax3, ax4

    @staticmethod
    def _subplot_pitch_notes(ax1, notes, pitch):
        # plot predominant melody
        ax1.plot(pitch[:, 0], pitch[:, 1], 'g', label='Pitch', alpha=0.7)

        # plot performed notes
        if notes is not None:
            Plotter._plot_performed_notes(ax1, notes)

        # axis style
        ax1.xaxis.set_label_coords(0.5, 0.05)
        ax1.set_xlabel('Time (sec)')
        ax1.set_ylabel('Frequency (Hz)')
        ax1.yaxis.grid(True)

        # move x-axis of ax1 and ax3 in between, e.g. top of ax1
        ax1.tick_params(axis='x', pad=-15)
        ax1.xaxis.set_label_position('top')

    @staticmethod
    def _plot_pitch_dist_note_models(ax2, note_models, pitch_distribution):
        # plot pitch distribution
        ax2.plot(pitch_distribution.vals, pitch_distribution.bins,
                 color='gray')

        # plot note models
        if note_models is not None:
            ytick_vals = Plotter._plot_note_models(
                ax2, note_models, pitch_distribution)
        else:
            peak_idx = pitch_distribution.detect_peaks()[0]
            ytick_vals = pitch_distribution.bins[peak_idx]

        # set the frequency ticks and grids
        ax2.set_yticks(ytick_vals)
        plt.setp(ax2.get_yticklabels(), visible=False)
        # ax2.yaxis.grid(True)

        # define xlim higher than the highest peak so the note names have space
        ax2.set_xlim([0, 1.2 * max(pitch_distribution.vals)])

        # remove spines from the second subplot
        ax2.spines['top'].set_visible(False)
        ax2.spines['bottom'].set_visible(False)
        ax2.spines['left'].set_visible(False)
        ax2.spines['right'].set_visible(False)

        # remove the axis of the subplot 2
        ax2.axis('off')

    @staticmethod
    def _plot_melodic_progression(ax3, melodic_progression, pitch,
                                  pitch_distribution):
        # plot...
        AudioSeyirAnalyzer.plot(melodic_progression, ax3)

        # axis style
        ax3.set_xlabel('')  # remove the automatically given labels
        ax3.set_ylabel('')
        plt.setp(ax3.get_yticklabels(), visible=False)
        plt.setp(ax3.get_xticklabels(), visible=False)

        # set xlim to the last time in the pitch track
        ax3.set_xlim([pitch[0, 0], pitch[-1, 0]])
        ax3.set_ylim([np.min(pitch_distribution.bins),
                      np.max(pitch_distribution.bins)])

        # remove the spines from the third subplot
        ax3.spines['bottom'].set_visible(False)
        ax3.spines['left'].set_visible(False)
        ax3.spines['right'].set_visible(False)
        ax3.get_yaxis().set_ticks([])

    @staticmethod
    def _plot_sections(ax1, ax3, ax4, sections):
        if sections is not None:
            sec_labels = []
            sec_locs = []
            xgrid_locs = []
            for sec in sections:
                # get the time interval
                tt = sec['time']
                dur = tt[1] - tt[0]
                xgrid_locs += tt  # add the boundaries to grid locations

                # get the plot limits
                ylim = ax3.get_ylim()

                # create the rectangle
                p = patches.Rectangle((tt[0], ylim[0]), dur, ylim[1],
                                      alpha=0.3)
                ax4.add_patch(p)

                sec_labels.append(sec['name'])
                sec_locs.append(np.mean(tt))

            # styling
            ax4.set_xticks(sec_locs)
            ax4.set_xticklabels(sec_labels, rotation=-15)

            ax1.set_xticks(xgrid_locs, minor=True)
            ax1.xaxis.grid(True, which='minor')
            ax1.xaxis.set_major_locator(FixedLocator(xgrid_locs,
                                                     nbins=len(xgrid_locs)/2))

            plt.setp(ax4.get_yticklabels(), visible=False)
            ax4.set_xlim(ax3.get_xlim())
        else:
            # no section labels
            plt.setp(ax4.get_xticklabels(), visible=False)

    @staticmethod
    def _plot_performed_notes(ax, notes):
        for note in notes:
            ax.plot(note['interval'], [note['performed_pitch']['value'],
                                       note['performed_pitch']['value']],
                    'r', alpha=0.4, linewidth=4)

    @staticmethod
    def _plot_note_models(ax2, note_models, pitch_distribution):
        max_rel_occur = Plotter._get_relative_note_occurences(
            note_models, pitch_distribution)

        Plotter._plot_note_distributions(ax2, note_models)

        ytick_vals = Plotter._plot_stable_pitches(
            ax2, max_rel_occur, note_models, max(pitch_distribution.vals))

        return ytick_vals

    @staticmethod
    def _plot_stable_pitches(ax, max_rel_occur, note_models, max_pd_height):
        ytick_vals = []
        for note_symbol, note in note_models.iteritems():
            if note['rel_occur'] > max_rel_occur * 0.1:
                ytick_vals.append(note['stable_pitch']['value'])

                # plot the performed frequency as a dashed line
                ax.hlines(y=note['theoretical_pitch']['value'], xmin=0,
                          xmax=note['rel_occur'], linestyles='dashed')

                # mark notes
                if note['performed_interval']['value'] == 0.0:  # tonic
                    ax.plot(note['rel_occur'], note['stable_pitch']['value'],
                            'cD', ms=10)
                else:
                    ax.plot(note['rel_occur'], note['stable_pitch']['value'],
                            'cD', ms=6, c='r')

                # print note name, lift the text a little bit
                txt_x_val = (note['rel_occur'] + 0.03 * max_pd_height)
                txt_str = ', '.join(
                    [note_symbol,
                     str(int(round(note['performed_interval']['value']))) +
                     ' cents'])
                ax.text(txt_x_val, note['stable_pitch']['value'], txt_str,
                        style='italic', horizontalalignment='left',
                        verticalalignment='center')
        return ytick_vals

    @staticmethod
    def _get_relative_note_occurences(note_models, pitch_distribution):
        max_rel_occur = 0
        for note_symbol, note in note_models.iteritems():
            # get the relative occurence of each note from the pitch
            # distribution
            dists = np.array([abs(note['stable_pitch']['value'] - dist_bin)
                              for dist_bin in pitch_distribution.bins])
            peak_ind = np.argmin(dists)
            note['rel_occur'] = pitch_distribution.vals[peak_ind]
            max_rel_occur = max([max_rel_occur, note['rel_occur']])
        return max_rel_occur

    @staticmethod
    def _plot_note_distributions(ax, note_models):
        for note_symbol, note in note_models.iteritems():
            try:
                ax.plot(note_models[note_symbol]['distribution'].vals,
                        note_models[note_symbol]['distribution'].bins,
                        label=note_symbol)
            except KeyError:
                pass  # note model is not available
<commit_msg>Section plotting (Closes #40)<commit_after>import numpy as np
from matplotlib import gridspec
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import FixedLocator
from seyiranalyzer.AudioSeyirAnalyzer import AudioSeyirAnalyzer

class Plotter(object):
    @staticmethod
    def plot_audio_features(pitch=None, pitch_distribution=None,
                            sections=None, notes=None, note_models=None,
                            melodic_progression=None):

        # create the figure and all the subplots with the shared axis specified
        fig, ax1, ax2, ax3, ax4 = Plotter._create_figure()

        # plot the pitch track and the performed notes to the first subplot
        Plotter._subplot_pitch_notes(ax1, notes, pitch)

        # plot the pitch distribution and the note models to the second subplot
        Plotter._plot_pitch_dist_note_models(ax2, note_models,
                                             pitch_distribution)

        # plot the melodic progression to the third subplot
        Plotter._plot_melodic_progression(ax3, melodic_progression, pitch,
                                          pitch_distribution)

        # plot the sections to the third subplot, onto the melodic progression
        Plotter._plot_sections(ax1, ax3, ax4, sections)

        return fig, (ax1, ax2, ax3, ax4)

    @staticmethod
    def _create_figure():
        # create the figure with four subplots with different size
        # - 1st is for the predominant melody and performed notes
        # - 2nd is the pitch distribution and note models, it shares the y
        # axis with the 1st
        # - 3rd is the melodic progression, it shares the x axis with the 1st
        # - 4th is for the sections, it is on top the 3rd
        fig = plt.figure()
        gs = gridspec.GridSpec(2, 2, width_ratios=[6, 1], height_ratios=[4, 1])
        ax1 = fig.add_subplot(gs[0])  # pitch and notes
        ax2 = fig.add_subplot(gs[1], sharey=ax1)  # pitch dist. and note models
        ax4 = fig.add_subplot(gs[2])  # sections
        ax3 = plt.twiny(ax4)  # melodic progression
        ax1.get_shared_x_axes().join(ax1, ax3)
        fig.subplots_adjust(hspace=0, wspace=0)
        return fig, ax1, ax2, ax3, ax4

    @staticmethod
    def _subplot_pitch_notes(ax1, notes, pitch):
        # plot predominant melody
        ax1.plot(pitch[:, 0], pitch[:, 1], 'g', label='Pitch', alpha=0.7)

        # plot performed notes
        if notes is not None:
            Plotter._plot_performed_notes(ax1, notes)

        # axis style
        ax1.xaxis.set_label_coords(0.5, 0.05)
        ax1.set_xlabel('Time (sec)')
        ax1.set_ylabel('Frequency (Hz)')
        ax1.yaxis.grid(True)

        # move x-axis of ax1 and ax3 in between, e.g. top of ax1
        ax1.tick_params(axis='x', pad=-15)
        ax1.xaxis.set_label_position('top')

    @staticmethod
    def _plot_pitch_dist_note_models(ax2, note_models, pitch_distribution):
        # plot pitch distribution
        ax2.plot(pitch_distribution.vals, pitch_distribution.bins,
                 color='gray')

        # plot note models
        if note_models is not None:
            ytick_vals = Plotter._plot_note_models(
                ax2, note_models, pitch_distribution)
        else:
            peak_idx = pitch_distribution.detect_peaks()[0]
            ytick_vals = pitch_distribution.bins[peak_idx]

        # set the frequency ticks and grids
        ax2.set_yticks(ytick_vals)
        plt.setp(ax2.get_yticklabels(), visible=False)
        # ax2.yaxis.grid(True)

        # define xlim higher than the highest peak so the note names have space
        ax2.set_xlim([0, 1.2 * max(pitch_distribution.vals)])

        # remove spines from the second subplot
        ax2.spines['top'].set_visible(False)
        ax2.spines['bottom'].set_visible(False)
        ax2.spines['left'].set_visible(False)
        ax2.spines['right'].set_visible(False)

        # remove the axis of the subplot 2
        ax2.axis('off')

    @staticmethod
    def _plot_melodic_progression(ax3, melodic_progression, pitch,
                                  pitch_distribution):
        # plot...
        AudioSeyirAnalyzer.plot(melodic_progression, ax3)

        # axis style
        ax3.set_xlabel('')  # remove the automatically given labels
        ax3.set_ylabel('')
        plt.setp(ax3.get_yticklabels(), visible=False)
        plt.setp(ax3.get_xticklabels(), visible=False)

        # set xlim to the last time in the pitch track
        ax3.set_xlim([pitch[0, 0], pitch[-1, 0]])
        ax3.set_ylim([np.min(pitch_distribution.bins),
                      np.max(pitch_distribution.bins)])

        # remove the spines from the third subplot
        ax3.spines['bottom'].set_visible(False)
        ax3.spines['left'].set_visible(False)
        ax3.spines['right'].set_visible(False)
        ax3.get_yaxis().set_ticks([])

    @staticmethod
    def _plot_sections(ax1, ax3, ax4, sections):
        if sections is not None:
            sec_labels = []
            sec_locs = []
            xgrid_locs = []

            unique_sections = list(set(s['name'] for s in sections))
            cmap = plt.get_cmap('gist_rainbow')
            colors = [cmap(i) for i in np.linspace(0, 1, len(unique_sections))]

            for sec in sections:
                # get the time interval
                tt = sec['time']
                dur = tt[1] - tt[0]
                xgrid_locs += tt  # add the boundaries to grid locations

                # get the plot limits
                ylim = ax3.get_ylim()

                # get the color code for the section
                clr = colors[unique_sections.index(sec['name'])]

                # create the rectangle
                p = patches.Rectangle((tt[0], ylim[0]), dur, ylim[1],
                                      alpha=0.3, color=clr)
                ax4.add_patch(p)

                sec_labels.append(sec['name'])
                sec_locs.append(np.mean(tt))

            # styling
            ax4.set_xticks(sec_locs)
            ax4.set_xticklabels(sec_labels, rotation=-15)

            ax1.set_xticks(xgrid_locs, minor=True)
            ax1.xaxis.grid(True, which='minor')
            ax1.xaxis.set_major_locator(FixedLocator(xgrid_locs,
                                                     nbins=len(xgrid_locs)/2))

            plt.setp(ax4.get_yticklabels(), visible=False)
            ax4.set_xlim(ax3.get_xlim())
        else:
            # no section labels
            plt.setp(ax4.get_xticklabels(), visible=False)

    @staticmethod
    def _plot_performed_notes(ax, notes):
        for note in notes:
            ax.plot(note['interval'], [note['performed_pitch']['value'],
                                       note['performed_pitch']['value']],
                    'r', alpha=0.4, linewidth=4)

    @staticmethod
    def _plot_note_models(ax2, note_models, pitch_distribution):
        max_rel_occur = Plotter._get_relative_note_occurences(
            note_models, pitch_distribution)

        Plotter._plot_note_distributions(ax2, note_models)

        ytick_vals = Plotter._plot_stable_pitches(
            ax2, max_rel_occur, note_models, max(pitch_distribution.vals))

        return ytick_vals

    @staticmethod
    def _plot_stable_pitches(ax, max_rel_occur, note_models, max_pd_height):
        ytick_vals = []
        for note_symbol, note in note_models.iteritems():
            if note['rel_occur'] > max_rel_occur * 0.1:
                ytick_vals.append(note['stable_pitch']['value'])

                # plot the performed frequency as a dashed line
                ax.hlines(y=note['theoretical_pitch']['value'], xmin=0,
                          xmax=note['rel_occur'], linestyles='dashed')

                # mark notes
                if note['performed_interval']['value'] == 0.0:  # tonic
                    ax.plot(note['rel_occur'], note['stable_pitch']['value'],
                            'cD', ms=10)
                else:
                    ax.plot(note['rel_occur'], note['stable_pitch']['value'],
                            'cD', ms=6, c='r')

                # print note name, lift the text a little bit
                txt_x_val = (note['rel_occur'] + 0.03 * max_pd_height)
                txt_str = ', '.join(
                    [note_symbol,
                     str(int(round(note['performed_interval']['value']))) +
                     ' cents'])
                ax.text(txt_x_val, note['stable_pitch']['value'], txt_str,
                        style='italic', horizontalalignment='left',
                        verticalalignment='center')
        return ytick_vals

    @staticmethod
    def _get_relative_note_occurences(note_models, pitch_distribution):
        max_rel_occur = 0
        for note_symbol, note in note_models.iteritems():
            # get the relative occurence of each note from the pitch
            # distribution
            dists = np.array([abs(note['stable_pitch']['value'] - dist_bin)
                              for dist_bin in pitch_distribution.bins])
            peak_ind = np.argmin(dists)
            note['rel_occur'] = pitch_distribution.vals[peak_ind]
            max_rel_occur = max([max_rel_occur, note['rel_occur']])
        return max_rel_occur

    @staticmethod
    def _plot_note_distributions(ax, note_models):
        for note_symbol, note in note_models.iteritems():
            try:
                ax.plot(note_models[note_symbol]['distribution'].vals,
                        note_models[note_symbol]['distribution'].bins,
                        label=note_symbol)
            except KeyError:
                pass  # note model is not available
<|endoftext|>


<commit_before># coding: utf-8

""" Print out a table of stars that would be good to observe now. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
from datetime import datetime
import warnings

# Third-party
from astropy import log as logger
from astropy.io import ascii
import astropy.coordinates as coord
import astropy.time as atime
import astropy.units as u
from astropy.utils.exceptions import AstropyWarning
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline

from gary.observation.rrlyrae import time_to_phase

kitt_peak = coord.EarthLocation.from_geodetic(lon=-111.5967*u.deg,
                                              lat=31.9583*u.deg,
                                              height=2000.*u.m)

def Vmag_to_exptime(V):
    # estimate exposure times
    s = InterpolatedUnivariateSpline([14.5,16,17.5], np.log10([150,600,2400]), k=1)
    y = 10**s(V)
    return y

def main(filename):
    filepath = os.path.abspath(os.path.expanduser(filename))
    tbl = ascii.read(filepath)
    ntargets = len(tbl)

    # get current time
    time = atime.Time(datetime.utcnow(), scale='utc')

    phases = np.zeros(ntargets)
    for i in range(ntargets):
        t0 = atime.Time(tbl[i]['mjd0'], scale='utc', format='mjd')
        phases[i] = time_to_phase(time, tbl[i]['period']*u.day, t0)

        # HACK: this corrects for the fact that CSS defines phase = 0.5
        #   to be peak brightness...Brani uses phase = 0.
        phases[i] = (phases[i] - 0.5) % 1.

    ix = (phases > 0.05) & (phases < 0.8)
    good_targets = tbl[ix]

    # add a column with estimated exposure times
    good_targets['exptime'] = [Vmag_to_exptime(V) for V in good_targets['VmagAvg']]

    # now add a "goodness" column
    g = coord.Galactic(l=good_targets['l']*u.deg,
                       b=good_targets['b']*u.deg)
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', AstropyWarning)
        altaz = g.transform_to(coord.AltAz(obstime=time, location=kitt_peak))
    airmass = altaz.secz
    good_targets['airmass'] = airmass

    # filter on airmass
    good_targets = good_targets[good_targets['airmass'] < 1.5]

    print('{0:<16} {1:<10} {2:<10} {3:<10}'.format("ID","V [mag]","airmass","exptime [sec]"))
    print('-'*(16+10+10+10+6))
    fmt_string = '{0:<16} {1:<10.2f} {2:<10.2f} {3:<10.0f}'
    for row in good_targets[['ID2015','VmagAvg','airmass','exptime']]:
        # print(row) #['ID2015','VmagAvg','airmass','exptime'])
        print(fmt_string.format(*row))

if __name__ == "__main__":
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")

    parser.add_argument("-f", "--file", dest="file", required=True,
                        type=str, help="Catalog file.")
    # parser.add_argument("-p", dest="plot", action="store_true", default=False,
    #                     help="Plot or not")

    args = parser.parse_args()

    # Set logger level based on verbose flags
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.ERROR)
    else:
        logger.setLevel(logging.INFO)

    main(filename=args.file)
<commit_msg>remove phase hacl<commit_after># coding: utf-8

""" Print out a table of stars that would be good to observe now. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
from datetime import datetime
import warnings

# Third-party
from astropy import log as logger
from astropy.io import ascii
import astropy.coordinates as coord
import astropy.time as atime
import astropy.units as u
from astropy.utils.exceptions import AstropyWarning
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline

from gary.observation.rrlyrae import time_to_phase

kitt_peak = coord.EarthLocation.from_geodetic(lon=-111.5967*u.deg,
                                              lat=31.9583*u.deg,
                                              height=2000.*u.m)

def Vmag_to_exptime(V):
    # estimate exposure times
    s = InterpolatedUnivariateSpline([14.5,16,17.5], np.log10([150,600,2400]), k=1)
    y = 10**s(V)
    return y

def main(filename):
    filepath = os.path.abspath(os.path.expanduser(filename))
    tbl = ascii.read(filepath)
    ntargets = len(tbl)

    # get current time
    time = atime.Time(datetime.utcnow(), scale='utc')

    phases = np.zeros(ntargets)
    for i in range(ntargets):
        t0 = atime.Time(tbl[i]['mjd0'], scale='utc', format='mjd')
        phases[i] = time_to_phase(time, tbl[i]['period']*u.day, t0)

        # HACK: this corrects for the fact that CSS defines phase = 0.5
        #   to be peak brightness...Brani uses phase = 0.
        # phases[i] = (phases[i] - 0.5) % 1.
        # UPDATE: actually, I think CSS uses the peak...

    ix = (phases > 0.05) & (phases < 0.8)
    good_targets = tbl[ix]

    # add a column with estimated exposure times
    good_targets['exptime'] = [Vmag_to_exptime(V) for V in good_targets['VmagAvg']]

    # now add a "goodness" column
    g = coord.Galactic(l=good_targets['l']*u.deg,
                       b=good_targets['b']*u.deg)
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', AstropyWarning)
        altaz = g.transform_to(coord.AltAz(obstime=time, location=kitt_peak))
    airmass = altaz.secz
    good_targets['airmass'] = airmass

    # filter on airmass
    good_targets = good_targets[good_targets['airmass'] < 1.5]

    print('{0:<16} {1:<10} {2:<10} {3:<10}'.format("ID","V [mag]","airmass","exptime [sec]"))
    print('-'*(16+10+10+10+6))
    fmt_string = '{0:<16} {1:<10.2f} {2:<10.2f} {3:<10.0f}'
    for row in good_targets[['ID2015','VmagAvg','airmass','exptime']]:
        print(fmt_string.format(*row))

if __name__ == "__main__":
    from argparse import ArgumentParser
    import logging

    # Define parser object
    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true", dest="verbose",
                        default=False, help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                        default=False, help="Be quiet! (default = False)")

    parser.add_argument("-f", "--file", dest="file", required=True,
                        type=str, help="Catalog file.")
    # parser.add_argument("-p", dest="plot", action="store_true", default=False,
    #                     help="Plot or not")

    args = parser.parse_args()

    # Set logger level based on verbose flags
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    elif args.quiet:
        logger.setLevel(logging.ERROR)
    else:
        logger.setLevel(logging.INFO)

    main(filename=args.file)
<|endoftext|>


<commit_before>import contextlib
import errno
import os
import sys
import tempfile

__version__ = '0.1.4'

PY2 = sys.version_info[0] == 2

text_type = unicode if PY2 else str

def _path_to_unicode(x):
    if not isinstance(x, text_type):
        return x.decode(sys.getfilesystemencoding())
    return x

if sys.platform != 'win32':
    def _replace_atomic(src, dst):
        os.rename(src, dst)

    def _move_atomic(src, dst):
        os.link(src, dst)
        os.unlink(src)
else:
    from ctypes import windll, WinError

    _MOVEFILE_REPLACE_EXISTING = 0x1
    _MOVEFILE_WRITE_THROUGH = 0x8
    _windows_default_flags = _MOVEFILE_WRITE_THROUGH

    def _handle_errors(rv):
        if not rv:
            raise WinError()

    def _replace_atomic(src, dst):
        _handle_errors(windll.kernel32.MoveFileExW(
            _path_to_unicode(src), _path_to_unicode(dst),
            _windows_default_flags | _MOVEFILE_REPLACE_EXISTING
        ))

    def _move_atomic(src, dst):
        _handle_errors(windll.kernel32.MoveFileExW(
            _path_to_unicode(src), _path_to_unicode(dst),
            _windows_default_flags
        ))

def replace_atomic(src, dst):
    '''
    Move ``src`` to ``dst``. If ``dst`` exists, it will be silently
    overwritten.

    Both paths must reside on the same filesystem for the operation to be
    atomic.
    '''
    return _replace_atomic(src, dst)

def move_atomic(src, dst):
    '''
    Move ``src`` to ``dst``. There might a timewindow where both filesystem
    entries exist. If ``dst`` already exists, :py:exc:`FileExistsError` will be
    raised.

    Both paths must reside on the same filesystem for the operation to be
    atomic.
    '''
    return _move_atomic(src, dst)

class AtomicWriter(object):
    '''
    A helper class for performing atomic writes. Usage::

        with AtomicWriter(path).open() as f:
            f.write(...)

    :param path: The destination filepath. May or may not exist.
    :param mode: The filemode for the temporary file.
    :param overwrite: If set to false, an error is raised if ``path`` exists.
        Errors are only raised after the file has been written to.  Either way,
        the operation is atomic.

    If you need further control over the exact behavior, you are encouraged to
    subclass.
    '''

    def __init__(self, path, mode='w', overwrite=False):
        if 'a' in mode:
            raise ValueError(
                'Appending to an existing file is not supported, because that '
                'would involve an expensive `copy`-operation to a temporary '
                'file. Open the file in normal `w`-mode and copy explicitly '
                'if that\'s what you\'re after.'
            )
        if 'x' in mode:
            raise ValueError('Use the `overwrite`-parameter instead.')
        if 'w' not in mode:
            raise ValueError('AtomicWriters can only be written to.')

        self._path = path
        self._mode = mode
        self._overwrite = overwrite

    def open(self):
        '''
        Open the temporary file.
        '''
        return self._open(self.get_fileobject)

    @contextlib.contextmanager
    def _open(self, get_fileobject):
        try:
            with get_fileobject() as f:
                yield f
            self.commit(f)
        except:
            try:
                self.rollback(f)
            except Exception:
                pass
            raise

    def get_fileobject(self, dir=None, **kwargs):
        '''Return the temporary path to use.'''
        if dir is None:
            dir = os.path.dirname(self._path)
        return tempfile.NamedTemporaryFile(mode=self._mode, dir=dir,
                                           delete=False, **kwargs)

    def commit(self, f):
        '''Move the temporary file to the target location.'''
        if self._overwrite:
            replace_atomic(f.name, self._path)  # atomic
        else:
            move_atomic(f.name, self._path)

    def rollback(self, f):
        '''Clean up all temporary resources.'''
        os.unlink(f.name)

def atomic_write(path, writer_cls=AtomicWriter, **cls_kwargs):
    '''
    Simple atomic writes. This wraps :py:class:`AtomicWriter`::

        with atomic_write(path) as f:
            f.write(...)

    :param path: The target path to write to.

    Additional keyword arguments are passed to the writer class. See
    :py:class:`AtomicWriterBase`.
    '''
    return writer_cls(path, **cls_kwargs).open()
<commit_msg>Version 0.1.5<commit_after>import contextlib
import errno
import os
import sys
import tempfile

__version__ = '0.1.5'

PY2 = sys.version_info[0] == 2

text_type = unicode if PY2 else str

def _path_to_unicode(x):
    if not isinstance(x, text_type):
        return x.decode(sys.getfilesystemencoding())
    return x

if sys.platform != 'win32':
    def _replace_atomic(src, dst):
        os.rename(src, dst)

    def _move_atomic(src, dst):
        os.link(src, dst)
        os.unlink(src)
else:
    from ctypes import windll, WinError

    _MOVEFILE_REPLACE_EXISTING = 0x1
    _MOVEFILE_WRITE_THROUGH = 0x8
    _windows_default_flags = _MOVEFILE_WRITE_THROUGH

    def _handle_errors(rv):
        if not rv:
            raise WinError()

    def _replace_atomic(src, dst):
        _handle_errors(windll.kernel32.MoveFileExW(
            _path_to_unicode(src), _path_to_unicode(dst),
            _windows_default_flags | _MOVEFILE_REPLACE_EXISTING
        ))

    def _move_atomic(src, dst):
        _handle_errors(windll.kernel32.MoveFileExW(
            _path_to_unicode(src), _path_to_unicode(dst),
            _windows_default_flags
        ))

def replace_atomic(src, dst):
    '''
    Move ``src`` to ``dst``. If ``dst`` exists, it will be silently
    overwritten.

    Both paths must reside on the same filesystem for the operation to be
    atomic.
    '''
    return _replace_atomic(src, dst)

def move_atomic(src, dst):
    '''
    Move ``src`` to ``dst``. There might a timewindow where both filesystem
    entries exist. If ``dst`` already exists, :py:exc:`FileExistsError` will be
    raised.

    Both paths must reside on the same filesystem for the operation to be
    atomic.
    '''
    return _move_atomic(src, dst)

class AtomicWriter(object):
    '''
    A helper class for performing atomic writes. Usage::

        with AtomicWriter(path).open() as f:
            f.write(...)

    :param path: The destination filepath. May or may not exist.
    :param mode: The filemode for the temporary file.
    :param overwrite: If set to false, an error is raised if ``path`` exists.
        Errors are only raised after the file has been written to.  Either way,
        the operation is atomic.

    If you need further control over the exact behavior, you are encouraged to
    subclass.
    '''

    def __init__(self, path, mode='w', overwrite=False):
        if 'a' in mode:
            raise ValueError(
                'Appending to an existing file is not supported, because that '
                'would involve an expensive `copy`-operation to a temporary '
                'file. Open the file in normal `w`-mode and copy explicitly '
                'if that\'s what you\'re after.'
            )
        if 'x' in mode:
            raise ValueError('Use the `overwrite`-parameter instead.')
        if 'w' not in mode:
            raise ValueError('AtomicWriters can only be written to.')

        self._path = path
        self._mode = mode
        self._overwrite = overwrite

    def open(self):
        '''
        Open the temporary file.
        '''
        return self._open(self.get_fileobject)

    @contextlib.contextmanager
    def _open(self, get_fileobject):
        try:
            with get_fileobject() as f:
                yield f
            self.commit(f)
        except:
            try:
                self.rollback(f)
            except Exception:
                pass
            raise

    def get_fileobject(self, dir=None, **kwargs):
        '''Return the temporary path to use.'''
        if dir is None:
            dir = os.path.dirname(self._path)
        return tempfile.NamedTemporaryFile(mode=self._mode, dir=dir,
                                           delete=False, **kwargs)

    def commit(self, f):
        '''Move the temporary file to the target location.'''
        if self._overwrite:
            replace_atomic(f.name, self._path)  # atomic
        else:
            move_atomic(f.name, self._path)

    def rollback(self, f):
        '''Clean up all temporary resources.'''
        os.unlink(f.name)

def atomic_write(path, writer_cls=AtomicWriter, **cls_kwargs):
    '''
    Simple atomic writes. This wraps :py:class:`AtomicWriter`::

        with atomic_write(path) as f:
            f.write(...)

    :param path: The target path to write to.

    Additional keyword arguments are passed to the writer class. See
    :py:class:`AtomicWriterBase`.
    '''
    return writer_cls(path, **cls_kwargs).open()
<|endoftext|>