<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Common functions and classes that are required by all query classes.
"""

import re
import warnings
import os
import shutil
import socket

import requests

from astropy.extern.six.moves.urllib_error import URLError
from astropy.extern import six
import astropy.units as u
from astropy import coordinates as coord
from collections import OrderedDict
import astropy.utils.data as aud
from astropy.io import fits, votable

from astropy.coordinates import BaseCoordinateFrame

from ..exceptions import TimeoutError
from .. import version

def ICRSCoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='icrs', **kwargs)

def GalacticCoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='galactic', **kwargs)

def FK5CoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='fk5', **kwargs)

def FK4CoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='fk4', **kwargs)

ICRSCoord = coord.SkyCoord
CoordClasses = (coord.SkyCoord, BaseCoordinateFrame)

__all__ = ['send_request',
           'parse_coordinates',
           'TableList',
           'suppress_vo_warnings',
           'validate_email']

def send_request(url, data, timeout, request_type='POST', headers={},
                 **kwargs):
    """
    A utility function that post HTTP requests to remote server
    and returns the HTTP response.

    Parameters
    ----------
    url : str
        The URL of the remote server
    data : dict
        A dictionary representing the payload to be posted via the HTTP request
    timeout : int, quantity_like
        Time limit for establishing successful connection with remote server
    request_type : str
        options are 'POST' (default) and 'GET'. Determines whether to perform
        an HTTP POST or an HTTP GET request
    headers : dict
        POST or GET headers.  user-agent will be set to
        astropy:astroquery.version

    Returns
    -------
    response : `requests.Response`
        Response object returned by the remote server
    """
    headers['User-Agent'] = ('astropy:astroquery.{vers}'
                             .format(vers=version.version))

    if hasattr(timeout, "unit"):
        warnings.warn("Converting timeout to seconds and truncating "
                      "to integer.")
        timeout = int(timeout.to(u.s).value)

    try:
        if request_type == 'GET':
            response = requests.get(url, params=data, timeout=timeout,
                                    headers=headers, **kwargs)
        elif request_type == 'POST':
            response = requests.post(url, data=data, timeout=timeout,
                                     headers=headers, **kwargs)
        else:
            raise ValueError("request_type must be either 'GET' or 'POST'.")

        response.raise_for_status()

        return response

    except requests.exceptions.Timeout:
        raise TimeoutError("Query timed out, time elapsed {time}s".
                           format(time=timeout))
    except requests.exceptions.RequestException as ex:
        raise Exception("Query failed: {0}\n".format(ex))

def radius_to_unit(radius, unit='degree'):
    """
    Helper function: Parse a radius, then return its value in degrees

    Parameters
    ----------
    radius : str or `~astropy.units.Quantity`
        The radius of a region

    Returns
    -------
    Floating point scalar value of radius in degrees
    """
    rad = coord.Angle(radius)

    if isinstance(unit, six.string_types):
        if hasattr(rad, unit):
            return getattr(rad, unit)
        elif hasattr(rad, unit + 's'):
            return getattr(rad, unit + 's')

    return rad.to(unit).value

def parse_coordinates(coordinates):
    """
    Takes a string or astropy.coordinates object. Checks if the
    string is parsable as an `astropy.coordinates`
    object or is a name that is resolvable. Otherwise asserts
    that the argument is an astropy.coordinates object.

    Parameters
    ----------
    coordinates : str or `astropy.coordinates` object
        Astronomical coordinate

    Returns
    -------
    coordinates : a subclass of `astropy.coordinates.BaseCoordinateFrame`

    Raises
    ------
    astropy.units.UnitsError
    TypeError
    """
    if isinstance(coordinates, six.string_types):
        try:
            c = ICRSCoord.from_name(coordinates)
        except coord.name_resolve.NameResolveError:
            try:
                c = ICRSCoordGenerator(coordinates)
                warnings.warn("Coordinate string is being interpreted as an "
                              "ICRS coordinate.")
            except u.UnitsError:
                warnings.warn("Only ICRS coordinates can be entered as "
                              "strings.\n For other systems please use the "
                              "appropriate astropy.coordinates object.")
                raise u.UnitsError
    elif isinstance(coordinates, CoordClasses):
        if hasattr(coordinates, 'frame'):
            c = coordinates
        else:
            # Convert the "frame" object into a SkyCoord
            c = coord.SkyCoord(coordinates)
    else:
        raise TypeError("Argument cannot be parsed as a coordinate")
    return c

def coord_to_radec(coordinate):
    """
    Wrapper to turn any astropy coordinate into FK5 RA in Hours and FK5 Dec in
    degrees

    This is a hack / temporary wrapper to deal with the unstable astropy API
    (it may be wise to remove this hack since it's not clear that the old
    coordinate API can even do transforms)
    """
    C = coordinate.transform_to('fk5')
    if hasattr(C.ra, 'hour'):
        ra = C.ra.hour
    elif hasattr(C.ra, 'hourangle'):
        ra = C.ra.hourangle
    else:
        raise ValueError("API Error: RA cannot be converted to hour "
                         "or hourangle.")
    dec = C.dec.degree
    return ra, dec

class TableList(list):

    """
    A class that inherits from `list` but included some pretty printing methods
    for an OrderedDict of `astropy.table.Table` objects.

    HINT: To access the tables by # instead of by table ID:
    >>> t = TableList([('a',1),('b',2)])
    >>> t[1]
    2
    >>> t['b']
    2
    """

    def __init__(self, inp):
        if not isinstance(inp, OrderedDict):
            try:
                inp = OrderedDict(inp)
            except (TypeError, ValueError):
                raise ValueError("Input to TableList must be an OrderedDict "
                                 "or list of (k,v) pairs")

        self._dict = inp
        super(TableList, self).__init__(inp.values())

    def __getitem__(self, key):
        if isinstance(key, int):
            # get the value in the (key,value) pair
            return super(TableList, self).__getitem__(key)
        elif key in self._dict:
            return self._dict[key]
        else:
            raise TypeError("TableLists can only be indexed with the "
                            "named keys and integers.")

    def __setitem__(self, value):
        raise TypeError("TableList is immutable.")

    def __getslice__(self, slice):
        return list(self.values())[slice]

    def keys(self):
        return list(self._dict.keys())

    def values(self):
        return list(self._dict.values())

    def __repr__(self):
        """
        Overrides the `OrderedDict.__repr__` method to return a simple summary
        of the `TableList` object.
        """

        return self.format_table_list()

    def format_table_list(self):
        """
        Prints the names of all `astropy.table.Table` objects, with their
        respective number of row and columns, contained in the
        `TableList` instance.
        """
        ntables = len(list(self.keys()))
        if ntables == 0:
            return "Empty TableList"

        header_str = "TableList with {keylen} tables:".format(keylen=ntables)
        body_str = "\n".join(["\t'{t_number}:{t_name}' with {ncol} column(s) "
                              "and {nrow} row(s) "
                              .format(t_number=t_number, t_name=t_name,
                                      nrow=len(self[t_number]),
                                      ncol=len(self[t_number].colnames))
                              for t_number, t_name in enumerate(self.keys())])
        return "\n".join([header_str, body_str])

    def print_table_list(self):
        print(self.format_table_list())

    def pprint(self, **kwargs):
        """ Helper function to make API more similar to astropy.Tables """
        if kwargs != {}:
            warnings.warn("TableList is a container of astropy.Tables.")
        self.print_table_list()

def _is_coordinate(coordinates):
    """
    Returns `True` if coordinates can be parsed via `astropy.coordinates`
    and `False` otherwise.

    Parameters
    ----------
    coordinates : str or `astropy.coordinates` object
            The target around which to search. It may be specified as a
            string in which case it is resolved using online services or as
            the appropriate `astropy.coordinates` object. ICRS coordinates
            may also be entered as strings as specified in the
            `astropy.coordinates` module.

    Returns
    -------
    bool
    """
    if hasattr(coordinates, 'fk5'):
        # its coordinate-like enough
        return True
    try:
        ICRSCoordGenerator(coordinates)
        return True
    except ValueError:
        return False

def suppress_vo_warnings():
    """
    Suppresses all warnings of the class
    `astropy.io.votable.exceptions.VOWarning`.
    """
    warnings.filterwarnings("ignore", category=votable.exceptions.VOWarning)

def validate_email(email):
    """
    E-mail address validation.  Uses validate_email if available, else a simple
    regex that will let through some invalid e-mails but will catch the most
    common violators.
    """
    try:
        import validate_email
        return validate_email.validate_email(email)
    except ImportError:
        return bool(re.compile(r'^\S+@\S+\.\S+$').match(email))

class FileContainer(object):
    """
    A File Object container, meant to offer lazy access to downloaded FITS
    files.
    """

    def __init__(self, target, **kwargs):
        kwargs.setdefault('cache', True)
        self._target = target
        self._timeout = kwargs.get('remote_timeout', aud.conf.remote_timeout)
        if (os.path.splitext(target)[1] == '.fits' and not
                ('encoding' in kwargs and kwargs['encoding'] == 'binary')):
            warnings.warn("FITS files must be read as binaries; error is "
                          "likely.")
        self._readable_object = get_readable_fileobj(target, **kwargs)

    def get_fits(self):
        """
        Assuming the contained file is a FITS file, read it
        and return the file parsed as FITS HDUList
        """
        filedata = self.get_string()

        self._fits = fits.HDUList.fromstring(filedata)

        return self._fits

    def save_fits(self, savepath, link_cache='hard'):
        """
        Save a FITS file to savepath

        Parameters
        ----------
        savepath : str
            The full path to a FITS filename, e.g. "file.fits", or
            "/path/to/file.fits".
        link_cache : 'hard', 'sym', or False
            Try to create a hard or symbolic link to the astropy cached file?
            If the system is unable to create a hardlink, the file will be
            copied to the target location.
        """
        from warnings import warn

        self.get_fits()

        try:
            dldir, urlmapfn = aud._get_download_cache_locs()
        except (IOError, OSError) as e:
            msg = 'Remote data cache could not be accessed due to '
            estr = '' if len(e.args) < 1 else (': ' + str(e))
            warn(aud.CacheMissingWarning(msg + e.__class__.__name__ + estr))

        with _open_shelve(urlmapfn, True) as url2hash:
            if str(self._target) in url2hash:
                target = url2hash[str(self._target)]
            else:
                raise IOError("Cached file not found / does not exist.")

        if link_cache == 'hard':
            try:
                os.link(target, savepath)
            except (IOError, OSError, AttributeError) as e:
                shutil.copy(target, savepath)
        elif link_cache == 'sym':
            try:
                os.symlink(target, savepath)
            except AttributeError:
                raise OSError('Creating symlinks is not possible on this OS.')
        else:
            shutil.copy(target, savepath)

    def get_string(self):
        """
        Download the file as a string
        """
        if not hasattr(self, '_string'):
            try:
                with self._readable_object as f:
                    data = f.read()
                    self._string = data
            except URLError as e:
                if isinstance(e.reason, socket.timeout):
                    raise TimeoutError("Query timed out, time elapsed {t}s".
                                       format(t=self._timeout))
                else:
                    raise e

        return self._string

    def get_stringio(self):
        """
        Return the file as an io.StringIO object
        """
        s = self.get_string()
        # TODO: replace with six.BytesIO
        try:
            return six.BytesIO(s)
        except TypeError:
            return six.StringIO(s)

    def __repr__(self):
        if hasattr(self, '_fits'):
            return "Downloaded FITS file: " + self._fits.__repr__()
        else:
            return ("Downloaded object from URL {} with ID {}"
                    .format(self._target, id(self._readable_object)))

def get_readable_fileobj(*args, **kwargs):
    """
    Overload astropy's get_readable_fileobj so that we can safely monkeypatch
    it in astroquery without affecting astropy core functionality
    """
    return aud.get_readable_fileobj(*args, **kwargs)

def parse_votable(content):
    """
    Parse a votable in string format
    """
    tables = votable.parse(six.BytesIO(content), pedantic=False)
    return tables

def _open_shelve(shelffn, withclosing=False):
    """
    Opens a shelf file.  If `withclosing` is  True, it will be opened with closing,
    allowing use like:

        with _open_shelve('somefile',True) as s:
            ...
    """
    import shelve
    import contextlib

    shelf = shelve.open(shelffn, protocol=2)

    if withclosing:
        return contextlib.closing(shelf)
    else:
        return shelf
<commit_msg>fix pycodestyle issue<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Common functions and classes that are required by all query classes.
"""

import re
import warnings
import os
import shutil
import socket

import requests

from astropy.extern.six.moves.urllib_error import URLError
from astropy.extern import six
import astropy.units as u
from astropy import coordinates as coord
from collections import OrderedDict
import astropy.utils.data as aud
from astropy.io import fits, votable

from astropy.coordinates import BaseCoordinateFrame

from ..exceptions import TimeoutError
from .. import version

def ICRSCoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='icrs', **kwargs)

def GalacticCoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='galactic', **kwargs)

def FK5CoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='fk5', **kwargs)

def FK4CoordGenerator(*args, **kwargs):
    return coord.SkyCoord(*args, frame='fk4', **kwargs)

ICRSCoord = coord.SkyCoord
CoordClasses = (coord.SkyCoord, BaseCoordinateFrame)

__all__ = ['send_request',
           'parse_coordinates',
           'TableList',
           'suppress_vo_warnings',
           'validate_email']

def send_request(url, data, timeout, request_type='POST', headers={},
                 **kwargs):
    """
    A utility function that post HTTP requests to remote server
    and returns the HTTP response.

    Parameters
    ----------
    url : str
        The URL of the remote server
    data : dict
        A dictionary representing the payload to be posted via the HTTP request
    timeout : int, quantity_like
        Time limit for establishing successful connection with remote server
    request_type : str
        options are 'POST' (default) and 'GET'. Determines whether to perform
        an HTTP POST or an HTTP GET request
    headers : dict
        POST or GET headers.  user-agent will be set to
        astropy:astroquery.version

    Returns
    -------
    response : `requests.Response`
        Response object returned by the remote server
    """
    headers['User-Agent'] = ('astropy:astroquery.{vers}'
                             .format(vers=version.version))

    if hasattr(timeout, "unit"):
        warnings.warn("Converting timeout to seconds and truncating "
                      "to integer.")
        timeout = int(timeout.to(u.s).value)

    try:
        if request_type == 'GET':
            response = requests.get(url, params=data, timeout=timeout,
                                    headers=headers, **kwargs)
        elif request_type == 'POST':
            response = requests.post(url, data=data, timeout=timeout,
                                     headers=headers, **kwargs)
        else:
            raise ValueError("request_type must be either 'GET' or 'POST'.")

        response.raise_for_status()

        return response

    except requests.exceptions.Timeout:
        raise TimeoutError("Query timed out, time elapsed {time}s".
                           format(time=timeout))
    except requests.exceptions.RequestException as ex:
        raise Exception("Query failed: {0}\n".format(ex))

def radius_to_unit(radius, unit='degree'):
    """
    Helper function: Parse a radius, then return its value in degrees

    Parameters
    ----------
    radius : str or `~astropy.units.Quantity`
        The radius of a region

    Returns
    -------
    Floating point scalar value of radius in degrees
    """
    rad = coord.Angle(radius)

    if isinstance(unit, six.string_types):
        if hasattr(rad, unit):
            return getattr(rad, unit)
        elif hasattr(rad, unit + 's'):
            return getattr(rad, unit + 's')

    return rad.to(unit).value

def parse_coordinates(coordinates):
    """
    Takes a string or astropy.coordinates object. Checks if the
    string is parsable as an `astropy.coordinates`
    object or is a name that is resolvable. Otherwise asserts
    that the argument is an astropy.coordinates object.

    Parameters
    ----------
    coordinates : str or `astropy.coordinates` object
        Astronomical coordinate

    Returns
    -------
    coordinates : a subclass of `astropy.coordinates.BaseCoordinateFrame`

    Raises
    ------
    astropy.units.UnitsError
    TypeError
    """
    if isinstance(coordinates, six.string_types):
        try:
            c = ICRSCoord.from_name(coordinates)
        except coord.name_resolve.NameResolveError:
            try:
                c = ICRSCoordGenerator(coordinates)
                warnings.warn("Coordinate string is being interpreted as an "
                              "ICRS coordinate.")
            except u.UnitsError:
                warnings.warn("Only ICRS coordinates can be entered as "
                              "strings.\n For other systems please use the "
                              "appropriate astropy.coordinates object.")
                raise u.UnitsError
    elif isinstance(coordinates, CoordClasses):
        if hasattr(coordinates, 'frame'):
            c = coordinates
        else:
            # Convert the "frame" object into a SkyCoord
            c = coord.SkyCoord(coordinates)
    else:
        raise TypeError("Argument cannot be parsed as a coordinate")
    return c

def coord_to_radec(coordinate):
    """
    Wrapper to turn any astropy coordinate into FK5 RA in Hours and FK5 Dec in
    degrees

    This is a hack / temporary wrapper to deal with the unstable astropy API
    (it may be wise to remove this hack since it's not clear that the old
    coordinate API can even do transforms)
    """
    C = coordinate.transform_to('fk5')
    if hasattr(C.ra, 'hour'):
        ra = C.ra.hour
    elif hasattr(C.ra, 'hourangle'):
        ra = C.ra.hourangle
    else:
        raise ValueError("API Error: RA cannot be converted to hour "
                         "or hourangle.")
    dec = C.dec.degree
    return ra, dec

class TableList(list):

    """
    A class that inherits from `list` but included some pretty printing methods
    for an OrderedDict of `astropy.table.Table` objects.

    HINT: To access the tables by # instead of by table ID:
    >>> t = TableList([('a',1),('b',2)])
    >>> t[1]
    2
    >>> t['b']
    2
    """

    def __init__(self, inp):
        if not isinstance(inp, OrderedDict):
            try:
                inp = OrderedDict(inp)
            except (TypeError, ValueError):
                raise ValueError("Input to TableList must be an OrderedDict "
                                 "or list of (k,v) pairs")

        self._dict = inp
        super(TableList, self).__init__(inp.values())

    def __getitem__(self, key):
        if isinstance(key, int):
            # get the value in the (key,value) pair
            return super(TableList, self).__getitem__(key)
        elif key in self._dict:
            return self._dict[key]
        else:
            raise TypeError("TableLists can only be indexed with the "
                            "named keys and integers.")

    def __setitem__(self, value):
        raise TypeError("TableList is immutable.")

    def __getslice__(self, slice):
        return list(self.values())[slice]

    def keys(self):
        return list(self._dict.keys())

    def values(self):
        return list(self._dict.values())

    def __repr__(self):
        """
        Overrides the `OrderedDict.__repr__` method to return a simple summary
        of the `TableList` object.
        """

        return self.format_table_list()

    def format_table_list(self):
        """
        Prints the names of all `astropy.table.Table` objects, with their
        respective number of row and columns, contained in the
        `TableList` instance.
        """
        ntables = len(list(self.keys()))
        if ntables == 0:
            return "Empty TableList"

        header_str = "TableList with {keylen} tables:".format(keylen=ntables)
        body_str = "\n".join(["\t'{t_number}:{t_name}' with {ncol} column(s) "
                              "and {nrow} row(s) "
                              .format(t_number=t_number, t_name=t_name,
                                      nrow=len(self[t_number]),
                                      ncol=len(self[t_number].colnames))
                              for t_number, t_name in enumerate(self.keys())])
        return "\n".join([header_str, body_str])

    def print_table_list(self):
        print(self.format_table_list())

    def pprint(self, **kwargs):
        """ Helper function to make API more similar to astropy.Tables """
        if kwargs != {}:
            warnings.warn("TableList is a container of astropy.Tables.")
        self.print_table_list()

def _is_coordinate(coordinates):
    """
    Returns `True` if coordinates can be parsed via `astropy.coordinates`
    and `False` otherwise.

    Parameters
    ----------
    coordinates : str or `astropy.coordinates` object
            The target around which to search. It may be specified as a
            string in which case it is resolved using online services or as
            the appropriate `astropy.coordinates` object. ICRS coordinates
            may also be entered as strings as specified in the
            `astropy.coordinates` module.

    Returns
    -------
    bool
    """
    if hasattr(coordinates, 'fk5'):
        # its coordinate-like enough
        return True
    try:
        ICRSCoordGenerator(coordinates)
        return True
    except ValueError:
        return False

def suppress_vo_warnings():
    """
    Suppresses all warnings of the class
    `astropy.io.votable.exceptions.VOWarning`.
    """
    warnings.filterwarnings("ignore", category=votable.exceptions.VOWarning)

def validate_email(email):
    """
    E-mail address validation.  Uses validate_email if available, else a simple
    regex that will let through some invalid e-mails but will catch the most
    common violators.
    """
    try:
        import validate_email
        return validate_email.validate_email(email)
    except ImportError:
        return bool(re.compile(r'^\S+@\S+\.\S+$').match(email))

class FileContainer(object):
    """
    A File Object container, meant to offer lazy access to downloaded FITS
    files.
    """

    def __init__(self, target, **kwargs):
        kwargs.setdefault('cache', True)
        self._target = target
        self._timeout = kwargs.get('remote_timeout', aud.conf.remote_timeout)
        if (os.path.splitext(target)[1] == '.fits' and not
                ('encoding' in kwargs and kwargs['encoding'] == 'binary')):
            warnings.warn("FITS files must be read as binaries; error is "
                          "likely.")
        self._readable_object = get_readable_fileobj(target, **kwargs)

    def get_fits(self):
        """
        Assuming the contained file is a FITS file, read it
        and return the file parsed as FITS HDUList
        """
        filedata = self.get_string()

        self._fits = fits.HDUList.fromstring(filedata)

        return self._fits

    def save_fits(self, savepath, link_cache='hard'):
        """
        Save a FITS file to savepath

        Parameters
        ----------
        savepath : str
            The full path to a FITS filename, e.g. "file.fits", or
            "/path/to/file.fits".
        link_cache : 'hard', 'sym', or False
            Try to create a hard or symbolic link to the astropy cached file?
            If the system is unable to create a hardlink, the file will be
            copied to the target location.
        """
        from warnings import warn

        self.get_fits()

        try:
            dldir, urlmapfn = aud._get_download_cache_locs()
        except (IOError, OSError) as e:
            msg = 'Remote data cache could not be accessed due to '
            estr = '' if len(e.args) < 1 else (': ' + str(e))
            warn(aud.CacheMissingWarning(msg + e.__class__.__name__ + estr))

        with _open_shelve(urlmapfn, True) as url2hash:
            if str(self._target) in url2hash:
                target = url2hash[str(self._target)]
            else:
                raise IOError("Cached file not found / does not exist.")

        if link_cache == 'hard':
            try:
                os.link(target, savepath)
            except (IOError, OSError, AttributeError) as e:
                shutil.copy(target, savepath)
        elif link_cache == 'sym':
            try:
                os.symlink(target, savepath)
            except AttributeError:
                raise OSError('Creating symlinks is not possible on this OS.')
        else:
            shutil.copy(target, savepath)

    def get_string(self):
        """
        Download the file as a string
        """
        if not hasattr(self, '_string'):
            try:
                with self._readable_object as f:
                    data = f.read()
                    self._string = data
            except URLError as e:
                if isinstance(e.reason, socket.timeout):
                    raise TimeoutError("Query timed out, time elapsed {t}s".
                                       format(t=self._timeout))
                else:
                    raise e

        return self._string

    def get_stringio(self):
        """
        Return the file as an io.StringIO object
        """
        s = self.get_string()
        # TODO: replace with six.BytesIO
        try:
            return six.BytesIO(s)
        except TypeError:
            return six.StringIO(s)

    def __repr__(self):
        if hasattr(self, '_fits'):
            return "Downloaded FITS file: " + self._fits.__repr__()
        else:
            return ("Downloaded object from URL {} with ID {}"
                    .format(self._target, id(self._readable_object)))

def get_readable_fileobj(*args, **kwargs):
    """
    Overload astropy's get_readable_fileobj so that we can safely monkeypatch
    it in astroquery without affecting astropy core functionality
    """
    return aud.get_readable_fileobj(*args, **kwargs)

def parse_votable(content):
    """
    Parse a votable in string format
    """
    tables = votable.parse(six.BytesIO(content), pedantic=False)
    return tables

def _open_shelve(shelffn, withclosing=False):
    """
    Opens a shelf file.  If `withclosing` is  True, it will be opened with closing,
    allowing use like:

        with _open_shelve('somefile',True) as s:
            ...
    """
    import shelve
    import contextlib

    shelf = shelve.open(shelffn, protocol=2)

    if withclosing:
        return contextlib.closing(shelf)
    else:
        return shelf
<|endoftext|>


<commit_before># encoding: utf-8
"""Tests for IPython.core.ultratb
"""
import io
import logging
import sys
import os.path
from textwrap import dedent
import traceback
import unittest
from unittest import mock

import IPython.core.ultratb as ultratb
from IPython.core.ultratb import ColorTB, VerboseTB, find_recursion

from IPython.testing import tools as tt
from IPython.testing.decorators import onlyif_unicode_paths
from IPython.utils.syspathcontext import prepended_to_syspath
from IPython.utils.tempdir import TemporaryDirectory

file_1 = """1
2
3
def f():
  1/0
"""

file_2 = """def f():
  1/0
"""

def recursionlimit(frames):
    """
    decorator to set the recursion limit temporarily
    """

    def inner(test_function):
        def wrapper(*args, **kwargs):
            _orig_rec_limit = ultratb._FRAME_RECURSION_LIMIT
            ultratb._FRAME_RECURSION_LIMIT = 50

            rl = sys.getrecursionlimit()
            sys.setrecursionlimit(frames)
            try:
                return test_function(*args, **kwargs)
            finally:
                sys.setrecursionlimit(rl)
                ultratb._FRAME_RECURSION_LIMIT = _orig_rec_limit

        return wrapper

    return inner

class ChangedPyFileTest(unittest.TestCase):
    def test_changing_py_file(self):
        """Traceback produced if the line where the error occurred is missing?

        https://github.com/ipython/ipython/issues/1456
        """
        with TemporaryDirectory() as td:
            fname = os.path.join(td, "foo.py")
            with open(fname, "w") as f:
                f.write(file_1)

            with prepended_to_syspath(td):
                ip.run_cell("import foo")

            with tt.AssertPrints("ZeroDivisionError"):
                ip.run_cell("foo.f()")

            # Make the file shorter, so the line of the error is missing.
            with open(fname, "w") as f:
                f.write(file_2)

            # For some reason, this was failing on the *second* call after
            # changing the file, so we call f() twice.
            with tt.AssertNotPrints("Internal Python error", channel='stderr'):
                with tt.AssertPrints("ZeroDivisionError"):
                    ip.run_cell("foo.f()")
                with tt.AssertPrints("ZeroDivisionError"):
                    ip.run_cell("foo.f()")

iso_8859_5_file = u'''# coding: iso-8859-5

def fail():
    """дбИЖ"""
    1/0     # дбИЖ
'''

class NonAsciiTest(unittest.TestCase):
    @onlyif_unicode_paths
    def test_nonascii_path(self):
        # Non-ascii directory name as well.
        with TemporaryDirectory(suffix=u'é') as td:
            fname = os.path.join(td, u"fooé.py")
            with open(fname, "w") as f:
                f.write(file_1)

            with prepended_to_syspath(td):
                ip.run_cell("import foo")

            with tt.AssertPrints("ZeroDivisionError"):
                ip.run_cell("foo.f()")

    def test_iso8859_5(self):
        with TemporaryDirectory() as td:
            fname = os.path.join(td, 'dfghjkl.py')

            with io.open(fname, 'w', encoding='iso-8859-5') as f:
                f.write(iso_8859_5_file)

            with prepended_to_syspath(td):
                ip.run_cell("from dfghjkl import fail")

            with tt.AssertPrints("ZeroDivisionError"):
                with tt.AssertPrints(u'дбИЖ', suppress=False):
                    ip.run_cell('fail()')

    def test_nonascii_msg(self):
        cell = u"raise Exception('é')"
        expected = u"Exception('é')"
        ip.run_cell("%xmode plain")
        with tt.AssertPrints(expected):
            ip.run_cell(cell)

        ip.run_cell("%xmode verbose")
        with tt.AssertPrints(expected):
            ip.run_cell(cell)

        ip.run_cell("%xmode context")
        with tt.AssertPrints(expected):
            ip.run_cell(cell)

        ip.run_cell("%xmode minimal")
        with tt.AssertPrints(u"Exception: é"):
            ip.run_cell(cell)

        # Put this back into Context mode for later tests.
        ip.run_cell("%xmode context")

class NestedGenExprTestCase(unittest.TestCase):
    """
    Regression test for the following issues:
    https://github.com/ipython/ipython/issues/8293
    https://github.com/ipython/ipython/issues/8205
    """
    def test_nested_genexpr(self):
        code = dedent(
            """\
            class SpecificException(Exception):
                pass

            def foo(x):
                raise SpecificException("Success!")

            sum(sum(foo(x) for _ in [0]) for x in [0])
            """
        )
        with tt.AssertPrints('SpecificException: Success!', suppress=False):
            ip.run_cell(code)

indentationerror_file = """if True:
zoon()
"""

class IndentationErrorTest(unittest.TestCase):
    def test_indentationerror_shows_line(self):
        # See issue gh-2398
        with tt.AssertPrints("IndentationError"):
            with tt.AssertPrints("zoon()", suppress=False):
                ip.run_cell(indentationerror_file)

        with TemporaryDirectory() as td:
            fname = os.path.join(td, "foo.py")
            with open(fname, "w") as f:
                f.write(indentationerror_file)

            with tt.AssertPrints("IndentationError"):
                with tt.AssertPrints("zoon()", suppress=False):
                    ip.magic('run %s' % fname)

se_file_1 = """1
2
7/
"""

se_file_2 = """7/
"""

class SyntaxErrorTest(unittest.TestCase):
    def test_syntaxerror_without_lineno(self):
        with tt.AssertNotPrints("TypeError"):
            with tt.AssertPrints("line unknown"):
                ip.run_cell("raise SyntaxError()")

    def test_syntaxerror_no_stacktrace_at_compile_time(self):
        syntax_error_at_compile_time = """
def foo():
    ..
"""
        with tt.AssertPrints("SyntaxError"):
            ip.run_cell(syntax_error_at_compile_time)

        with tt.AssertNotPrints("foo()"):
            ip.run_cell(syntax_error_at_compile_time)

    def test_syntaxerror_stacktrace_when_running_compiled_code(self):
        syntax_error_at_runtime = """
def foo():
    eval("..")

def bar():
    foo()

bar()
"""
        with tt.AssertPrints("SyntaxError"):
            ip.run_cell(syntax_error_at_runtime)
        # Assert syntax error during runtime generate stacktrace
        with tt.AssertPrints(["foo()", "bar()"]):
            ip.run_cell(syntax_error_at_runtime)
        del ip.user_ns['bar']
        del ip.user_ns['foo']

    def test_changing_py_file(self):
        with TemporaryDirectory() as td:
            fname = os.path.join(td, "foo.py")
            with open(fname, 'w') as f:
                f.write(se_file_1)

            with tt.AssertPrints(["7/", "SyntaxError"]):
                ip.magic("run " + fname)

            # Modify the file
            with open(fname, 'w') as f:
                f.write(se_file_2)

            # The SyntaxError should point to the correct line
            with tt.AssertPrints(["7/", "SyntaxError"]):
                ip.magic("run " + fname)

    def test_non_syntaxerror(self):
        # SyntaxTB may be called with an error other than a SyntaxError
        # See e.g. gh-4361
        try:
            raise ValueError('QWERTY')
        except ValueError:
            with tt.AssertPrints('QWERTY'):
                ip.showsyntaxerror()

class Python3ChainedExceptionsTest(unittest.TestCase):
    DIRECT_CAUSE_ERROR_CODE = """
try:
    x = 1 + 2
    print(not_defined_here)
except Exception as e:
    x += 55
    x - 1
    y = {}
    raise KeyError('uh') from e
    """

    EXCEPTION_DURING_HANDLING_CODE = """
try:
    x = 1 + 2
    print(not_defined_here)
except Exception as e:
    x += 55
    x - 1
    y = {}
    raise KeyError('uh')
    """

    SUPPRESS_CHAINING_CODE = """
try:
    1/0
except Exception:
    raise ValueError("Yikes") from None
    """

    def test_direct_cause_error(self):
        with tt.AssertPrints(["KeyError", "NameError", "direct cause"]):
            ip.run_cell(self.DIRECT_CAUSE_ERROR_CODE)

    def test_exception_during_handling_error(self):
        with tt.AssertPrints(["KeyError", "NameError", "During handling"]):
            ip.run_cell(self.EXCEPTION_DURING_HANDLING_CODE)

    def test_suppress_exception_chaining(self):
        with tt.AssertNotPrints("ZeroDivisionError"), \
             tt.AssertPrints("ValueError", suppress=False):
            ip.run_cell(self.SUPPRESS_CHAINING_CODE)

    def test_plain_exception(self):
        with tt.AssertPrints(["KeyError", "NameError", "During handling"]):
            ip.run_cell("%xmode Plain")
            ip.run_cell(self.EXCEPTION_DURING_HANDLING_CODE)
            ip.run_cell("%xmode Verbose")

class RecursionTest(unittest.TestCase):
    DEFINITIONS = """
def non_recurs():
    1/0

def r1():
    r1()

def r3a():
    r3b()

def r3b():
    r3c()

def r3c():
    r3a()

def r3o1():
    r3a()

def r3o2():
    r3o1()
"""
    def setUp(self):
        ip.run_cell(self.DEFINITIONS)

    def test_no_recursion(self):
        with tt.AssertNotPrints("frames repeated"):
            ip.run_cell("non_recurs()")

    @recursionlimit(150)
    def test_recursion_one_frame(self):
        with tt.AssertPrints("1 frames repeated"):
            ip.run_cell("r1()")

    @recursionlimit(150)
    def test_recursion_three_frames(self):
        with tt.AssertPrints("3 frames repeated"):
            ip.run_cell("r3o2()")

    @recursionlimit(150)
    def test_find_recursion(self):
        captured = []
        def capture_exc(*args, **kwargs):
            captured.append(sys.exc_info())
        with mock.patch.object(ip, 'showtraceback', capture_exc):
            ip.run_cell("r3o2()")

        self.assertEqual(len(captured), 1)
        etype, evalue, tb = captured[0]
        self.assertIn("recursion", str(evalue))

        records = ip.InteractiveTB.get_records(tb, 3, ip.InteractiveTB.tb_offset)
        for r in records[:10]:
            print(r[1:4])

        # The outermost frames should be:
        # 0: the 'cell' that was running when the exception came up
        # 1: r3o2()
        # 2: r3o1()
        # 3: r3a()
        # Then repeating r3b, r3c, r3a
        last_unique, repeat_length = find_recursion(etype, evalue, records)
        self.assertEqual(last_unique, 2)
        self.assertEqual(repeat_length, 3)

#----------------------------------------------------------------------------

# module testing (minimal)
def test_handlers():
    def spam(c, d_e):
        (d, e) = d_e
        x = c + d
        y = c * d
        foo(x, y)

    def foo(a, b, bar=1):
        eggs(a, b + bar)

    def eggs(f, g, z=globals()):
        h = f + g
        i = f - g
        return h / i

    buff = io.StringIO()

    buff.write('')
    buff.write('*** Before ***')
    try:
        buff.write(spam(1, (2, 3)))
    except:
        traceback.print_exc(file=buff)

    handler = ColorTB(ostream=buff)
    buff.write('*** ColorTB ***')
    try:
        buff.write(spam(1, (2, 3)))
    except:
        handler(*sys.exc_info())
    buff.write('')

    handler = VerboseTB(ostream=buff)
    buff.write('*** VerboseTB ***')
    try:
        buff.write(spam(1, (2, 3)))
    except:
        handler(*sys.exc_info())
    buff.write('')

from IPython.testing.decorators import skipif

class TokenizeFailureTest(unittest.TestCase):
    """Tests related to https://github.com/ipython/ipython/issues/6864."""

    # that appear to test that we are handling an exception that can be thrown
    # by the tokenizer due to a bug that seem to have been fixed in 3.8, though
    # I'm unsure if other sequences can make it raise this error. Let's just
    # skip in 3.8 for now
    @skipif(sys.version_info > (3,8))
    def testLogging(self):
        message = "An unexpected error occurred while tokenizing input"
        cell = 'raise ValueError("""a\nb""")'

        stream = io.StringIO()
        handler = logging.StreamHandler(stream)
        logger = logging.getLogger()
        loglevel = logger.level
        logger.addHandler(handler)
        self.addCleanup(lambda: logger.removeHandler(handler))
        self.addCleanup(lambda: logger.setLevel(loglevel))

        logger.setLevel(logging.INFO)
        with tt.AssertNotPrints(message):
            ip.run_cell(cell)
        self.assertNotIn(message, stream.getvalue())

        logger.setLevel(logging.DEBUG)
        with tt.AssertNotPrints(message):
            ip.run_cell(cell)
        self.assertIn(message, stream.getvalue())
<commit_msg>Add more tests<commit_after># encoding: utf-8
"""Tests for IPython.core.ultratb
"""
import io
import logging
import sys
import os.path
from textwrap import dedent
import traceback
import unittest
from unittest import mock

import IPython.core.ultratb as ultratb
from IPython.core.ultratb import ColorTB, VerboseTB, find_recursion

from IPython.testing import tools as tt
from IPython.testing.decorators import onlyif_unicode_paths
from IPython.utils.syspathcontext import prepended_to_syspath
from IPython.utils.tempdir import TemporaryDirectory

file_1 = """1
2
3
def f():
  1/0
"""

file_2 = """def f():
  1/0
"""

def recursionlimit(frames):
    """
    decorator to set the recursion limit temporarily
    """

    def inner(test_function):
        def wrapper(*args, **kwargs):
            _orig_rec_limit = ultratb._FRAME_RECURSION_LIMIT
            ultratb._FRAME_RECURSION_LIMIT = 50

            rl = sys.getrecursionlimit()
            sys.setrecursionlimit(frames)
            try:
                return test_function(*args, **kwargs)
            finally:
                sys.setrecursionlimit(rl)
                ultratb._FRAME_RECURSION_LIMIT = _orig_rec_limit

        return wrapper

    return inner

class ChangedPyFileTest(unittest.TestCase):
    def test_changing_py_file(self):
        """Traceback produced if the line where the error occurred is missing?

        https://github.com/ipython/ipython/issues/1456
        """
        with TemporaryDirectory() as td:
            fname = os.path.join(td, "foo.py")
            with open(fname, "w") as f:
                f.write(file_1)

            with prepended_to_syspath(td):
                ip.run_cell("import foo")

            with tt.AssertPrints("ZeroDivisionError"):
                ip.run_cell("foo.f()")

            # Make the file shorter, so the line of the error is missing.
            with open(fname, "w") as f:
                f.write(file_2)

            # For some reason, this was failing on the *second* call after
            # changing the file, so we call f() twice.
            with tt.AssertNotPrints("Internal Python error", channel='stderr'):
                with tt.AssertPrints("ZeroDivisionError"):
                    ip.run_cell("foo.f()")
                with tt.AssertPrints("ZeroDivisionError"):
                    ip.run_cell("foo.f()")

iso_8859_5_file = u'''# coding: iso-8859-5

def fail():
    """дбИЖ"""
    1/0     # дбИЖ
'''

class NonAsciiTest(unittest.TestCase):
    @onlyif_unicode_paths
    def test_nonascii_path(self):
        # Non-ascii directory name as well.
        with TemporaryDirectory(suffix=u'é') as td:
            fname = os.path.join(td, u"fooé.py")
            with open(fname, "w") as f:
                f.write(file_1)

            with prepended_to_syspath(td):
                ip.run_cell("import foo")

            with tt.AssertPrints("ZeroDivisionError"):
                ip.run_cell("foo.f()")

    def test_iso8859_5(self):
        with TemporaryDirectory() as td:
            fname = os.path.join(td, 'dfghjkl.py')

            with io.open(fname, 'w', encoding='iso-8859-5') as f:
                f.write(iso_8859_5_file)

            with prepended_to_syspath(td):
                ip.run_cell("from dfghjkl import fail")

            with tt.AssertPrints("ZeroDivisionError"):
                with tt.AssertPrints(u'дбИЖ', suppress=False):
                    ip.run_cell('fail()')

    def test_nonascii_msg(self):
        cell = u"raise Exception('é')"
        expected = u"Exception('é')"
        ip.run_cell("%xmode plain")
        with tt.AssertPrints(expected):
            ip.run_cell(cell)

        ip.run_cell("%xmode verbose")
        with tt.AssertPrints(expected):
            ip.run_cell(cell)

        ip.run_cell("%xmode context")
        with tt.AssertPrints(expected):
            ip.run_cell(cell)

        ip.run_cell("%xmode minimal")
        with tt.AssertPrints(u"Exception: é"):
            ip.run_cell(cell)

        # Put this back into Context mode for later tests.
        ip.run_cell("%xmode context")

class NestedGenExprTestCase(unittest.TestCase):
    """
    Regression test for the following issues:
    https://github.com/ipython/ipython/issues/8293
    https://github.com/ipython/ipython/issues/8205
    """
    def test_nested_genexpr(self):
        code = dedent(
            """\
            class SpecificException(Exception):
                pass

            def foo(x):
                raise SpecificException("Success!")

            sum(sum(foo(x) for _ in [0]) for x in [0])
            """
        )
        with tt.AssertPrints('SpecificException: Success!', suppress=False):
            ip.run_cell(code)

indentationerror_file = """if True:
zoon()
"""

class IndentationErrorTest(unittest.TestCase):
    def test_indentationerror_shows_line(self):
        # See issue gh-2398
        with tt.AssertPrints("IndentationError"):
            with tt.AssertPrints("zoon()", suppress=False):
                ip.run_cell(indentationerror_file)

        with TemporaryDirectory() as td:
            fname = os.path.join(td, "foo.py")
            with open(fname, "w") as f:
                f.write(indentationerror_file)

            with tt.AssertPrints("IndentationError"):
                with tt.AssertPrints("zoon()", suppress=False):
                    ip.magic('run %s' % fname)

se_file_1 = """1
2
7/
"""

se_file_2 = """7/
"""

class SyntaxErrorTest(unittest.TestCase):
    def test_syntaxerror_without_lineno(self):
        with tt.AssertNotPrints("TypeError"):
            with tt.AssertPrints("line unknown"):
                ip.run_cell("raise SyntaxError()")

    def test_syntaxerror_no_stacktrace_at_compile_time(self):
        syntax_error_at_compile_time = """
def foo():
    ..
"""
        with tt.AssertPrints("SyntaxError"):
            ip.run_cell(syntax_error_at_compile_time)

        with tt.AssertNotPrints("foo()"):
            ip.run_cell(syntax_error_at_compile_time)

    def test_syntaxerror_stacktrace_when_running_compiled_code(self):
        syntax_error_at_runtime = """
def foo():
    eval("..")

def bar():
    foo()

bar()
"""
        with tt.AssertPrints("SyntaxError"):
            ip.run_cell(syntax_error_at_runtime)
        # Assert syntax error during runtime generate stacktrace
        with tt.AssertPrints(["foo()", "bar()"]):
            ip.run_cell(syntax_error_at_runtime)
        del ip.user_ns['bar']
        del ip.user_ns['foo']

    def test_changing_py_file(self):
        with TemporaryDirectory() as td:
            fname = os.path.join(td, "foo.py")
            with open(fname, 'w') as f:
                f.write(se_file_1)

            with tt.AssertPrints(["7/", "SyntaxError"]):
                ip.magic("run " + fname)

            # Modify the file
            with open(fname, 'w') as f:
                f.write(se_file_2)

            # The SyntaxError should point to the correct line
            with tt.AssertPrints(["7/", "SyntaxError"]):
                ip.magic("run " + fname)

    def test_non_syntaxerror(self):
        # SyntaxTB may be called with an error other than a SyntaxError
        # See e.g. gh-4361
        try:
            raise ValueError('QWERTY')
        except ValueError:
            with tt.AssertPrints('QWERTY'):
                ip.showsyntaxerror()

class Python3ChainedExceptionsTest(unittest.TestCase):
    DIRECT_CAUSE_ERROR_CODE = """
try:
    x = 1 + 2
    print(not_defined_here)
except Exception as e:
    x += 55
    x - 1
    y = {}
    raise KeyError('uh') from e
    """

    EXCEPTION_DURING_HANDLING_CODE = """
try:
    x = 1 + 2
    print(not_defined_here)
except Exception as e:
    x += 55
    x - 1
    y = {}
    raise KeyError('uh')
    """

    SUPPRESS_CHAINING_CODE = """
try:
    1/0
except Exception:
    raise ValueError("Yikes") from None
    """

    def test_direct_cause_error(self):
        with tt.AssertPrints(["KeyError", "NameError", "direct cause"]):
            ip.run_cell(self.DIRECT_CAUSE_ERROR_CODE)

    def test_exception_during_handling_error(self):
        with tt.AssertPrints(["KeyError", "NameError", "During handling"]):
            ip.run_cell(self.EXCEPTION_DURING_HANDLING_CODE)

    def test_suppress_exception_chaining(self):
        with tt.AssertNotPrints("ZeroDivisionError"), \
             tt.AssertPrints("ValueError", suppress=False):
            ip.run_cell(self.SUPPRESS_CHAINING_CODE)

    def test_plain_direct_cause_error(self):
        with tt.AssertPrints(["KeyError", "NameError", "direct cause"]):
            ip.run_cell("%xmode Plain")
            ip.run_cell(self.DIRECT_CAUSE_ERROR_CODE)
            ip.run_cell("%xmode Verbose")

    def test_plain_exception_during_handling_error(self):
        with tt.AssertPrints(["KeyError", "NameError", "During handling"]):
            ip.run_cell("%xmode Plain")
            ip.run_cell(self.EXCEPTION_DURING_HANDLING_CODE)
            ip.run_cell("%xmode Verbose")

    def test_plain_suppress_exception_chaining(self):
        with tt.AssertNotPrints("ZeroDivisionError"), \
             tt.AssertPrints("ValueError", suppress=False):
            ip.run_cell("%xmode Plain")
            ip.run_cell(self.SUPPRESS_CHAINING_CODE)
            ip.run_cell("%xmode Verbose")

class RecursionTest(unittest.TestCase):
    DEFINITIONS = """
def non_recurs():
    1/0

def r1():
    r1()

def r3a():
    r3b()

def r3b():
    r3c()

def r3c():
    r3a()

def r3o1():
    r3a()

def r3o2():
    r3o1()
"""
    def setUp(self):
        ip.run_cell(self.DEFINITIONS)

    def test_no_recursion(self):
        with tt.AssertNotPrints("frames repeated"):
            ip.run_cell("non_recurs()")

    @recursionlimit(150)
    def test_recursion_one_frame(self):
        with tt.AssertPrints("1 frames repeated"):
            ip.run_cell("r1()")

    @recursionlimit(150)
    def test_recursion_three_frames(self):
        with tt.AssertPrints("3 frames repeated"):
            ip.run_cell("r3o2()")

    @recursionlimit(150)
    def test_find_recursion(self):
        captured = []
        def capture_exc(*args, **kwargs):
            captured.append(sys.exc_info())
        with mock.patch.object(ip, 'showtraceback', capture_exc):
            ip.run_cell("r3o2()")

        self.assertEqual(len(captured), 1)
        etype, evalue, tb = captured[0]
        self.assertIn("recursion", str(evalue))

        records = ip.InteractiveTB.get_records(tb, 3, ip.InteractiveTB.tb_offset)
        for r in records[:10]:
            print(r[1:4])

        # The outermost frames should be:
        # 0: the 'cell' that was running when the exception came up
        # 1: r3o2()
        # 2: r3o1()
        # 3: r3a()
        # Then repeating r3b, r3c, r3a
        last_unique, repeat_length = find_recursion(etype, evalue, records)
        self.assertEqual(last_unique, 2)
        self.assertEqual(repeat_length, 3)

#----------------------------------------------------------------------------

# module testing (minimal)
def test_handlers():
    def spam(c, d_e):
        (d, e) = d_e
        x = c + d
        y = c * d
        foo(x, y)

    def foo(a, b, bar=1):
        eggs(a, b + bar)

    def eggs(f, g, z=globals()):
        h = f + g
        i = f - g
        return h / i

    buff = io.StringIO()

    buff.write('')
    buff.write('*** Before ***')
    try:
        buff.write(spam(1, (2, 3)))
    except:
        traceback.print_exc(file=buff)

    handler = ColorTB(ostream=buff)
    buff.write('*** ColorTB ***')
    try:
        buff.write(spam(1, (2, 3)))
    except:
        handler(*sys.exc_info())
    buff.write('')

    handler = VerboseTB(ostream=buff)
    buff.write('*** VerboseTB ***')
    try:
        buff.write(spam(1, (2, 3)))
    except:
        handler(*sys.exc_info())
    buff.write('')

from IPython.testing.decorators import skipif

class TokenizeFailureTest(unittest.TestCase):
    """Tests related to https://github.com/ipython/ipython/issues/6864."""

    # that appear to test that we are handling an exception that can be thrown
    # by the tokenizer due to a bug that seem to have been fixed in 3.8, though
    # I'm unsure if other sequences can make it raise this error. Let's just
    # skip in 3.8 for now
    @skipif(sys.version_info > (3,8))
    def testLogging(self):
        message = "An unexpected error occurred while tokenizing input"
        cell = 'raise ValueError("""a\nb""")'

        stream = io.StringIO()
        handler = logging.StreamHandler(stream)
        logger = logging.getLogger()
        loglevel = logger.level
        logger.addHandler(handler)
        self.addCleanup(lambda: logger.removeHandler(handler))
        self.addCleanup(lambda: logger.setLevel(loglevel))

        logger.setLevel(logging.INFO)
        with tt.AssertNotPrints(message):
            ip.run_cell(cell)
        self.assertNotIn(message, stream.getvalue())

        logger.setLevel(logging.DEBUG)
        with tt.AssertNotPrints(message):
            ip.run_cell(cell)
        self.assertIn(message, stream.getvalue())
<|endoftext|>


<commit_before># 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 httpretty

from keystoneclient.auth import token_endpoint
from keystoneclient import session
from keystoneclient.tests import utils

class TokenEndpointTest(utils.TestCase):

    TEST_TOKEN = 'aToken'
    TEST_URL = 'http://server/prefix'

    @httpretty.activate
    def test_basic_case(self):
        httpretty.register_uri(httpretty.GET, self.TEST_URL, body='body')

        a = token_endpoint.Token(self.TEST_URL, self.TEST_TOKEN)
        s = session.Session(auth=a)

        data = s.get(self.TEST_URL, authenticated=True)

        self.assertEqual(data.text, 'body')
        self.assertRequestHeaderEqual('X-Auth-Token', self.TEST_TOKEN)
<commit_msg>Add endpoint handling to Token/Endpoint auth<commit_after># 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 httpretty

from keystoneclient.auth import token_endpoint
from keystoneclient import session
from keystoneclient.tests import utils

class TokenEndpointTest(utils.TestCase):

    TEST_TOKEN = 'aToken'
    TEST_URL = 'http://server/prefix'

    @httpretty.activate
    def test_basic_case(self):
        httpretty.register_uri(httpretty.GET, self.TEST_URL, body='body')

        a = token_endpoint.Token(self.TEST_URL, self.TEST_TOKEN)
        s = session.Session(auth=a)

        data = s.get(self.TEST_URL, authenticated=True)

        self.assertEqual(data.text, 'body')
        self.assertRequestHeaderEqual('X-Auth-Token', self.TEST_TOKEN)

    @httpretty.activate
    def test_basic_endpoint_case(self):
        self.stub_url(httpretty.GET, ['p'], body='body')
        a = token_endpoint.Token(self.TEST_URL, self.TEST_TOKEN)
        s = session.Session(auth=a)

        data = s.get('/p',
                     authenticated=True,
                     endpoint_filter={'service': 'identity'})

        self.assertEqual(self.TEST_URL, a.get_endpoint(s))
        self.assertEqual('body', data.text)
        self.assertRequestHeaderEqual('X-Auth-Token', self.TEST_TOKEN)
<|endoftext|>


<commit_before>"""toplevel setup/teardown for parallel tests."""

#-------------------------------------------------------------------------------
#  Copyright (C) 2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-------------------------------------------------------------------------------

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

import os
import tempfile
import time
from subprocess import Popen, PIPE, STDOUT

from IPython.utils.path import get_ipython_dir
from IPython.parallel import Client

processes = []
blackhole = tempfile.TemporaryFile()

# nose setup/teardown

def setup():
    cp = Popen('ipcontroller --profile iptest -r --log-level 10 --log-to-file'.split(), stdout=blackhole, stderr=STDOUT)
    engine_json = os.path.join(get_ipython_dir(), 'cluster_iptest', 'security', 'ipcontroller-engine.json')
    client_json = os.path.join(get_ipython_dir(), 'cluster_iptest', 'security', 'ipcontroller-client.json')
    while not os.path.exists(engine_json) and not os.path.exists(client_json):
        time.sleep(0.1)
    processes.append(cp)
    add_engines(1)
    c = Client(profile='iptest')
    while not c.ids:
        time.sleep(.1)
        c.spin()
    c.close()

def add_engines(n=1, profile='iptest'):
    rc = Client(profile=profile)
    base = len(rc)
    eps = []
    for i in range(n):
        ep = Popen(['ipengine']+ ['--profile', profile, '--log-level', '10', '--log-to-file'], stdout=blackhole, stderr=STDOUT)
        # ep.start()
        processes.append(ep)
        eps.append(ep)
    while len(rc) < base+n:
        time.sleep(.1)
        rc.spin()
    rc.close()
    return eps

def teardown():
    time.sleep(1)
    while processes:
        p = processes.pop()
        if p.poll() is None:
            try:
                p.terminate()
            except Exception, e:
                print e
                pass
        if p.poll() is None:
            time.sleep(.25)
        if p.poll() is None:
            try:
                print 'killing'
                p.kill()
            except:
                print "couldn't shutdown process: ", p

<commit_msg>better handle controller/engines never starting in tests<commit_after>"""toplevel setup/teardown for parallel tests."""

#-------------------------------------------------------------------------------
#  Copyright (C) 2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-------------------------------------------------------------------------------

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

import os
import tempfile
import time
from subprocess import Popen, PIPE, STDOUT

from IPython.utils.path import get_ipython_dir
from IPython.parallel import Client

processes = []
blackhole = tempfile.TemporaryFile()

# nose setup/teardown

def setup():
    cp = Popen('ipcontroller --profile iptest -r --log-level 10 --log-to-file --usethreads'.split(), stdout=blackhole, stderr=STDOUT)
    processes.append(cp)
    engine_json = os.path.join(get_ipython_dir(), 'cluster_iptest', 'security', 'ipcontroller-engine.json')
    client_json = os.path.join(get_ipython_dir(), 'cluster_iptest', 'security', 'ipcontroller-client.json')
    tic = time.time()
    while not os.path.exists(engine_json) or not os.path.exists(client_json):
        if cp.poll() is not None:
            print cp.poll()
            raise RuntimeError("The test controller failed to start.")
        elif time.time()-tic > 10:
            raise RuntimeError("Timeout waiting for the test controller to start.")
        time.sleep(0.1)
    add_engines(1)

def add_engines(n=1, profile='iptest'):
    rc = Client(profile=profile)
    base = len(rc)
    eps = []
    for i in range(n):
        ep = Popen(['ipengine']+ ['--profile', profile, '--log-level', '10', '--log-to-file'], stdout=blackhole, stderr=STDOUT)
        # ep.start()
        processes.append(ep)
        eps.append(ep)
    tic = time.time()
    while len(rc) < base+n:
        if any([ ep.poll() is not None for ep in eps ]):
            raise RuntimeError("A test engine failed to start.")
        elif time.time()-tic > 10:
            raise RuntimeError("Timeout waiting for engines to connect.")
        time.sleep(.1)
        rc.spin()
    rc.close()
    return eps

def teardown():
    time.sleep(1)
    while processes:
        p = processes.pop()
        if p.poll() is None:
            try:
                p.terminate()
            except Exception, e:
                print e
                pass
        if p.poll() is None:
            time.sleep(.25)
        if p.poll() is None:
            try:
                print 'killing'
                p.kill()
            except:
                print "couldn't shutdown process: ", p

<|endoftext|>


<commit_before>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Various names for properties, transforms, etc."""

# All constants are for internal use only; no backwards-compatibility
# guarantees.

# pytype: skip-file

from __future__ import absolute_import

# Standard file names used for staging files.
from builtins import object

# Referenced by Dataflow legacy worker.
from apache_beam.runners.internal.names import PICKLED_MAIN_SESSION_FILE  # pylint: disable=unused-import

# String constants related to sources framework
SOURCE_FORMAT = 'custom_source'
SOURCE_TYPE = 'CustomSourcesType'
SERIALIZED_SOURCE_KEY = 'serialized_source'

# In a released SDK, container tags are selected based on the SDK version.
# Unreleased versions use container versions based on values of
# BEAM_CONTAINER_VERSION and BEAM_FNAPI_CONTAINER_VERSION (see below).

# Update this version to the next version whenever there is a change that will
# require changes to legacy Dataflow worker execution environment.
BEAM_CONTAINER_VERSION = 'beam-master-20201005'
# Update this version to the next version whenever there is a change that
# requires changes to SDK harness container or SDK harness launcher.
BEAM_FNAPI_CONTAINER_VERSION = 'beam-master-20201005'

DATAFLOW_CONTAINER_IMAGE_REPOSITORY = 'gcr.io/cloud-dataflow/v1beta3'

class TransformNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Transform strings as they are expected in the CloudWorkflow protos.
  """
  COLLECTION_TO_SINGLETON = 'CollectionToSingleton'
  COMBINE = 'CombineValues'
  CREATE_PCOLLECTION = 'CreateCollection'
  DO = 'ParallelDo'
  FLATTEN = 'Flatten'
  GROUP = 'GroupByKey'
  READ = 'ParallelRead'
  WRITE = 'ParallelWrite'

class PropertyNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Property strings as they are expected in the CloudWorkflow protos.
  """
  BIGQUERY_CREATE_DISPOSITION = 'create_disposition'
  BIGQUERY_DATASET = 'dataset'
  BIGQUERY_EXPORT_FORMAT = 'bigquery_export_format'
  BIGQUERY_FLATTEN_RESULTS = 'bigquery_flatten_results'
  BIGQUERY_KMS_KEY = 'bigquery_kms_key'
  BIGQUERY_PROJECT = 'project'
  BIGQUERY_QUERY = 'bigquery_query'
  BIGQUERY_SCHEMA = 'schema'
  BIGQUERY_TABLE = 'table'
  BIGQUERY_USE_LEGACY_SQL = 'bigquery_use_legacy_sql'
  BIGQUERY_WRITE_DISPOSITION = 'write_disposition'
  DISPLAY_DATA = 'display_data'
  ELEMENT = 'element'
  ELEMENTS = 'elements'
  ENCODING = 'encoding'
  FILE_PATTERN = 'filepattern'
  FILE_NAME_PREFIX = 'filename_prefix'
  FILE_NAME_SUFFIX = 'filename_suffix'
  FORMAT = 'format'
  INPUTS = 'inputs'
  IMPULSE_ELEMENT = 'impulse_element'
  NON_PARALLEL_INPUTS = 'non_parallel_inputs'
  NUM_SHARDS = 'num_shards'
  OUT = 'out'
  OUTPUT = 'output'
  OUTPUT_INFO = 'output_info'
  OUTPUT_NAME = 'output_name'
  PARALLEL_INPUT = 'parallel_input'
  PIPELINE_PROTO_TRANSFORM_ID = 'pipeline_proto_transform_id'
  PUBSUB_ID_LABEL = 'pubsub_id_label'
  PUBSUB_SERIALIZED_ATTRIBUTES_FN = 'pubsub_serialized_attributes_fn'
  PUBSUB_SUBSCRIPTION = 'pubsub_subscription'
  PUBSUB_TIMESTAMP_ATTRIBUTE = 'pubsub_timestamp_label'
  PUBSUB_TOPIC = 'pubsub_topic'
  RESTRICTION_ENCODING = 'restriction_encoding'
  SERIALIZED_FN = 'serialized_fn'
  SHARD_NAME_TEMPLATE = 'shard_template'
  SOURCE_STEP_INPUT = 'custom_source_step_input'
  SERIALIZED_TEST_STREAM = 'serialized_test_stream'
  STEP_NAME = 'step_name'
  USE_INDEXED_FORMAT = 'use_indexed_format'
  USER_FN = 'user_fn'
  USER_NAME = 'user_name'
  USES_KEYED_STATE = 'uses_keyed_state'
  VALIDATE_SINK = 'validate_sink'
  VALIDATE_SOURCE = 'validate_source'
  VALUE = 'value'
  WINDOWING_STRATEGY = 'windowing_strategy'
<commit_msg>Update Beam Dataflow container versions for Python<commit_after>#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Various names for properties, transforms, etc."""

# All constants are for internal use only; no backwards-compatibility
# guarantees.

# pytype: skip-file

from __future__ import absolute_import

# Standard file names used for staging files.
from builtins import object

# Referenced by Dataflow legacy worker.
from apache_beam.runners.internal.names import PICKLED_MAIN_SESSION_FILE  # pylint: disable=unused-import

# String constants related to sources framework
SOURCE_FORMAT = 'custom_source'
SOURCE_TYPE = 'CustomSourcesType'
SERIALIZED_SOURCE_KEY = 'serialized_source'

# In a released SDK, container tags are selected based on the SDK version.
# Unreleased versions use container versions based on values of
# BEAM_CONTAINER_VERSION and BEAM_FNAPI_CONTAINER_VERSION (see below).

# Update this version to the next version whenever there is a change that will
# require changes to legacy Dataflow worker execution environment.
BEAM_CONTAINER_VERSION = 'beam-master-20201116'
# Update this version to the next version whenever there is a change that
# requires changes to SDK harness container or SDK harness launcher.
BEAM_FNAPI_CONTAINER_VERSION = 'beam-master-20201116'

DATAFLOW_CONTAINER_IMAGE_REPOSITORY = 'gcr.io/cloud-dataflow/v1beta3'

class TransformNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Transform strings as they are expected in the CloudWorkflow protos.
  """
  COLLECTION_TO_SINGLETON = 'CollectionToSingleton'
  COMBINE = 'CombineValues'
  CREATE_PCOLLECTION = 'CreateCollection'
  DO = 'ParallelDo'
  FLATTEN = 'Flatten'
  GROUP = 'GroupByKey'
  READ = 'ParallelRead'
  WRITE = 'ParallelWrite'

class PropertyNames(object):
  """For internal use only; no backwards-compatibility guarantees.

  Property strings as they are expected in the CloudWorkflow protos.
  """
  BIGQUERY_CREATE_DISPOSITION = 'create_disposition'
  BIGQUERY_DATASET = 'dataset'
  BIGQUERY_EXPORT_FORMAT = 'bigquery_export_format'
  BIGQUERY_FLATTEN_RESULTS = 'bigquery_flatten_results'
  BIGQUERY_KMS_KEY = 'bigquery_kms_key'
  BIGQUERY_PROJECT = 'project'
  BIGQUERY_QUERY = 'bigquery_query'
  BIGQUERY_SCHEMA = 'schema'
  BIGQUERY_TABLE = 'table'
  BIGQUERY_USE_LEGACY_SQL = 'bigquery_use_legacy_sql'
  BIGQUERY_WRITE_DISPOSITION = 'write_disposition'
  DISPLAY_DATA = 'display_data'
  ELEMENT = 'element'
  ELEMENTS = 'elements'
  ENCODING = 'encoding'
  FILE_PATTERN = 'filepattern'
  FILE_NAME_PREFIX = 'filename_prefix'
  FILE_NAME_SUFFIX = 'filename_suffix'
  FORMAT = 'format'
  INPUTS = 'inputs'
  IMPULSE_ELEMENT = 'impulse_element'
  NON_PARALLEL_INPUTS = 'non_parallel_inputs'
  NUM_SHARDS = 'num_shards'
  OUT = 'out'
  OUTPUT = 'output'
  OUTPUT_INFO = 'output_info'
  OUTPUT_NAME = 'output_name'
  PARALLEL_INPUT = 'parallel_input'
  PIPELINE_PROTO_TRANSFORM_ID = 'pipeline_proto_transform_id'
  PUBSUB_ID_LABEL = 'pubsub_id_label'
  PUBSUB_SERIALIZED_ATTRIBUTES_FN = 'pubsub_serialized_attributes_fn'
  PUBSUB_SUBSCRIPTION = 'pubsub_subscription'
  PUBSUB_TIMESTAMP_ATTRIBUTE = 'pubsub_timestamp_label'
  PUBSUB_TOPIC = 'pubsub_topic'
  RESTRICTION_ENCODING = 'restriction_encoding'
  SERIALIZED_FN = 'serialized_fn'
  SHARD_NAME_TEMPLATE = 'shard_template'
  SOURCE_STEP_INPUT = 'custom_source_step_input'
  SERIALIZED_TEST_STREAM = 'serialized_test_stream'
  STEP_NAME = 'step_name'
  USE_INDEXED_FORMAT = 'use_indexed_format'
  USER_FN = 'user_fn'
  USER_NAME = 'user_name'
  USES_KEYED_STATE = 'uses_keyed_state'
  VALIDATE_SINK = 'validate_sink'
  VALIDATE_SOURCE = 'validate_source'
  VALUE = 'value'
  WINDOWING_STRATEGY = 'windowing_strategy'
<|endoftext|>