<commit_before>from Queue import Queue
import subprocess
import threading
import traceback
import logging
import numpy
import time

log = logging.getLogger(__name__)

"""
	Quick and dirty, frame-aware MP3 encoding bridge using LAME.
	About 75% of the speed of raw LAME. Pass PCM data to the Lame class,
	get back (via callback, queue or file) MP3 frames. Supports real-time
	encoding or blocking for the length of the audio stream - useful for
	an MP3 server, or something else real time, for example.
"""

"""
Some important LAME facts used below:
	Each MP3 frame is identifiable by a header.
	This header has, essentially:
		"Frame Sync"            11 1's (i.e.: 0xFF + 3 bits)
		"Mpeg Audio Version ID" should be 0b11 for MPEG V1, 0b10 for MPEG V2
		"Layer Description"     should be 0b11
		"Protection Bit"        set to 1 by Lame, not protected
		"Bitrate index"         0000 -> free
								0001 -> 32 kbps
								0010 -> 40 kbps
								0011 -> 48 kbps
								0100 -> 56 kbps
								0101 -> 64 kbps
								0110 -> 80 kbps
								0111 -> 96 kbps
								1000 -> 112 kbps
								1001 -> 128 kbps
								1010 -> 160 kbps
								1011 -> 192 kbps
								1100 -> 224 kbps
								1101 -> 256 kbps
								1110 -> 320 kbps
								1111 -> invalid

	Following the header, there are always SAMPLES_PER_FRAME samples of audio data.
	At our constant sampling frequency of 44100, this means each frame
	contains exactly .026122449 seconds of audio.
"""

BITRATE_TABLE = [
	0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, None
]
SAMPLERATE_TABLE = [
	44100, 48000, 32000, None
]
HEADER_SIZE = 4
SAMPLES_PER_FRAME = 1152

def avg(l):
	return sum(l) / len(l)

def frame_length(header):
	bitrate = BITRATE_TABLE[ord(header[2]) >> 4]
	sample_rate = SAMPLERATE_TABLE[(ord(header[2]) & 0b00001100) >> 2]
	padding = (ord(header[2]) & 0b00000010) >> 1
	return int((float(SAMPLES_PER_FRAME) / sample_rate)
			   * ((bitrate / 8) * 1000)) + padding

def synchsafe(n):
	"""Return four bytes as synchsafe encoded version of the 28-bit number n."""
	# TODO: Simplify/optimize this.
	return str(bytearray([n>>21,(n>>14)&127,(n>>7)&127,n&127]))

class Lame(threading.Thread):
	"""
		Live MP3 streamer. Currently only works for 16-bit, 44.1kHz stereo input.
	"""
	safety_buffer = 30  # seconds
	input_wordlength = 16
	samplerate = 44100
	channels = 2
	preset = "-V3"

	#   Time-sensitive options
	real_time = False       #   Should we encode in 1:1 real time?
	block = False           #   Regardless of real-time, should we block
							#   for as long as the audio we've encoded lasts?

	stream_chunk_size = samplerate / 8
	data = None

	def __init__(self, callback=None, ofile=None, oqueue=None, syncqueue=None):
		threading.Thread.__init__(self)

		self.lame = None
		self.buffered = 0
		self.lame_input_length = 0
		self.in_samples = 0
		self.out_samples = 0
		self.delta = 0
		self.oqueue = oqueue
		self.syncqueue = syncqueue
		self.ofile = ofile
		self.callback = callback

		self.markers = []
		self.finished = False
		self.sent = False
		self.ready = threading.Semaphore()
		self.encode = threading.Semaphore()
		self.setDaemon(True)

		self.__write_queue = Queue()
		self.__write_thread = threading.Thread(target=self.__lame_write)
		self.__write_thread.setDaemon(True)
		self.__write_thread.start()

	@property
	def pcm_datarate(self):
		return self.samplerate * self.channels * (self.input_wordlength / 8)

	def add_pcm(self, data, marker=None):
		"""
		Expects PCM data in the form of a NumPy array,
		or an AudioRenderable that will be sliced according to the start and end.

		"""
		if self.lame.returncode is not None:
			return False
		if marker: self.markers.append((self.in_samples, marker))
		self.encode.acquire()
		if isinstance(data, numpy.ndarray):
			samples = len(data)
		else:
			samples = data.samples
		self.__write_queue.put(data)
		self.in_samples += samples
		put_time = time.time()
		if self.buffered >= self.safety_buffer:
			self.ready.acquire()
		done_time = time.time()
		if self.block and not self.real_time:
			delay = (samples / float(self.samplerate)) \
					- (done_time - put_time) \
					- self.safety_buffer
			time.sleep(delay)
		return True

	def __lame_write(self):
		while not self.finished:
			try:
				data = self.__write_queue.get()
				if data is None:
					break
				if isinstance(data, numpy.ndarray):
					self.buffered += len(data) / self.channels \
											* (self.input_wordlength / 8)
					try:
						data.tofile(self.lame.stdin)
					except IOError:
						log.error("Could not write to lame!")
						self.finished = True
						break
				else:
					try:
						tmp = 0
						for chunk in data.render(self.stream_chunk_size):
							try:
								samples = len(chunk)
								self.buffered += samples
								self.lame_input_length += samples
								tmp += samples
								chunk.tofile(self.lame.stdin)
							except IOError:
								log.error("Could not write to lame!")
								self.finished = True
								break
						self.delta += tmp - data.samples
						log.debug("Current delta: %d samples.", self.delta)
						#   Note: this delta will cause drift of 1 second/month.
						#   TODO: Fix it. Eventually.
					except Exception:
						log.error("Couldn't render segment due to:\n%s",
								traceback.format_exc())
			except Exception:
				log.critical("Failed to write to Lame:\n%s",
							 traceback.format_exc())
			finally:
				self.encode.release()
		log.critical("Encoder finishing!")

	#   TODO: Extend me to work for all samplerates
	def start(self, *args, **kwargs):
		call = ["lame"]
		call.append('-r')
		if self.input_wordlength != 16:
			call.extend(["--bitwidth", str(self.input_wordlength)])
		call.extend(self.preset.split())
		call.extend(["-", "-"])
		self.lame = subprocess.Popen(
			call,
			stdin=subprocess.PIPE,
			stdout=subprocess.PIPE,
			stderr=subprocess.PIPE
		)
		threading.Thread.start(self, *args, **kwargs)

	def ensure_is_alive(self):
		if self.finished:
			return False
		if self.is_alive():
			return True
		try:
			self.start()
			return True
		except Exception:
			return False

	def run(self, *args, **kwargs):
		try:
			last = None
			lag = 0
			while True:
				timing = float(SAMPLES_PER_FRAME) / self.samplerate

				header = self.lame.stdout.read(HEADER_SIZE)
				if len(header) == HEADER_SIZE:
					frame_len = frame_length(header) - HEADER_SIZE
					frame = self.lame.stdout.read(frame_len)
					buf = header + frame
					if len(frame) == frame_len:
						self.buffered -= SAMPLES_PER_FRAME
				else:
					buf = header

				if self.buffered < (self.safety_buffer * self.samplerate):
					self.ready.release()
				if len(buf):
					self.out_samples += SAMPLES_PER_FRAME
					if self.ofile:
						self.ofile.write(buf)
						self.ofile.flush()
					if self.callback:
						self.callback(False)
					if self.markers and self.markers[0][0] <= self.out_samples:
						msg = self.markers.pop(0)[1];
						log.info("Adding marker: %r", msg)
						if isinstance(msg, unicode): msg=msg.encode("utf-8")
						msg = "\3" + msg # Prepend the UTF-8 encoding marker. If msg was already bytes, it MUST be valid UTF-8.
						frame = "TIT2"+synchsafe(len(msg))+"\0\0"+msg;
						self.oqueue.put("ID3\4\0\0"+synchsafe(len(frame))+frame)
					if self.oqueue:
						self.oqueue.put(buf)
					if self.real_time and self.sent:
						now = time.time()
						if last:
							delta = (now - last - timing)
							lag += delta
							if lag < timing:
								time.sleep(max(0, timing - delta))
						last = now
					self.sent = True
				else:
					if self.callback:
						self.callback(True)
					break
			self.lame.wait()
		except Exception:
			log.error(traceback.format_exc())
			self.finish()
			raise

	def finish(self):
		"""
			Closes input stream to LAME and waits for the last frame(s) to
			finish encoding. Returns LAME's return value code.
		"""
		if self.lame:
			self.__write_queue.put(None)
			self.encode.acquire()
			self.lame.stdin.close()
			self.join()
			self.finished = True
			return self.lame.returncode
		return -1

if __name__ == "__main__":
	import wave
	f = wave.open("test.wav")
	a = numpy.frombuffer(f.readframes(f.getnframes()),
						 dtype=numpy.int16).reshape((-1, 2))

	s = time.time()
	print "Encoding test.wav to testout.mp3..."
	encoder = Lame(ofile=open('testout.mp3', 'w'))
	encoder.safety_buffer = 30
	encoder.start()
	encoder.add_pcm(a)
	encoder.finish()
	s = time.time() - s
	print "Took %2.2fs" % s
<commit_msg>Downgrade ID3 version used<commit_after>from Queue import Queue
import subprocess
import threading
import traceback
import logging
import numpy
import time

log = logging.getLogger(__name__)

"""
	Quick and dirty, frame-aware MP3 encoding bridge using LAME.
	About 75% of the speed of raw LAME. Pass PCM data to the Lame class,
	get back (via callback, queue or file) MP3 frames. Supports real-time
	encoding or blocking for the length of the audio stream - useful for
	an MP3 server, or something else real time, for example.
"""

"""
Some important LAME facts used below:
	Each MP3 frame is identifiable by a header.
	This header has, essentially:
		"Frame Sync"            11 1's (i.e.: 0xFF + 3 bits)
		"Mpeg Audio Version ID" should be 0b11 for MPEG V1, 0b10 for MPEG V2
		"Layer Description"     should be 0b11
		"Protection Bit"        set to 1 by Lame, not protected
		"Bitrate index"         0000 -> free
								0001 -> 32 kbps
								0010 -> 40 kbps
								0011 -> 48 kbps
								0100 -> 56 kbps
								0101 -> 64 kbps
								0110 -> 80 kbps
								0111 -> 96 kbps
								1000 -> 112 kbps
								1001 -> 128 kbps
								1010 -> 160 kbps
								1011 -> 192 kbps
								1100 -> 224 kbps
								1101 -> 256 kbps
								1110 -> 320 kbps
								1111 -> invalid

	Following the header, there are always SAMPLES_PER_FRAME samples of audio data.
	At our constant sampling frequency of 44100, this means each frame
	contains exactly .026122449 seconds of audio.
"""

BITRATE_TABLE = [
	0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, None
]
SAMPLERATE_TABLE = [
	44100, 48000, 32000, None
]
HEADER_SIZE = 4
SAMPLES_PER_FRAME = 1152

def avg(l):
	return sum(l) / len(l)

def frame_length(header):
	bitrate = BITRATE_TABLE[ord(header[2]) >> 4]
	sample_rate = SAMPLERATE_TABLE[(ord(header[2]) & 0b00001100) >> 2]
	padding = (ord(header[2]) & 0b00000010) >> 1
	return int((float(SAMPLES_PER_FRAME) / sample_rate)
			   * ((bitrate / 8) * 1000)) + padding

def synchsafe(n):
	"""Return four bytes as synchsafe encoded version of the 28-bit number n."""
	# TODO: Simplify/optimize this.
	return str(bytearray([n>>21,(n>>14)&127,(n>>7)&127,n&127]))

class Lame(threading.Thread):
	"""
		Live MP3 streamer. Currently only works for 16-bit, 44.1kHz stereo input.
	"""
	safety_buffer = 30  # seconds
	input_wordlength = 16
	samplerate = 44100
	channels = 2
	preset = "-V3"

	#   Time-sensitive options
	real_time = False       #   Should we encode in 1:1 real time?
	block = False           #   Regardless of real-time, should we block
							#   for as long as the audio we've encoded lasts?

	stream_chunk_size = samplerate / 8
	data = None

	def __init__(self, callback=None, ofile=None, oqueue=None, syncqueue=None):
		threading.Thread.__init__(self)

		self.lame = None
		self.buffered = 0
		self.lame_input_length = 0
		self.in_samples = 0
		self.out_samples = 0
		self.delta = 0
		self.oqueue = oqueue
		self.syncqueue = syncqueue
		self.ofile = ofile
		self.callback = callback

		self.markers = []
		self.finished = False
		self.sent = False
		self.ready = threading.Semaphore()
		self.encode = threading.Semaphore()
		self.setDaemon(True)

		self.__write_queue = Queue()
		self.__write_thread = threading.Thread(target=self.__lame_write)
		self.__write_thread.setDaemon(True)
		self.__write_thread.start()

	@property
	def pcm_datarate(self):
		return self.samplerate * self.channels * (self.input_wordlength / 8)

	def add_pcm(self, data, marker=None):
		"""
		Expects PCM data in the form of a NumPy array,
		or an AudioRenderable that will be sliced according to the start and end.

		"""
		if self.lame.returncode is not None:
			return False
		if marker: self.markers.append((self.in_samples, marker))
		self.encode.acquire()
		if isinstance(data, numpy.ndarray):
			samples = len(data)
		else:
			samples = data.samples
		self.__write_queue.put(data)
		self.in_samples += samples
		put_time = time.time()
		if self.buffered >= self.safety_buffer:
			self.ready.acquire()
		done_time = time.time()
		if self.block and not self.real_time:
			delay = (samples / float(self.samplerate)) \
					- (done_time - put_time) \
					- self.safety_buffer
			time.sleep(delay)
		return True

	def __lame_write(self):
		while not self.finished:
			try:
				data = self.__write_queue.get()
				if data is None:
					break
				if isinstance(data, numpy.ndarray):
					self.buffered += len(data) / self.channels \
											* (self.input_wordlength / 8)
					try:
						data.tofile(self.lame.stdin)
					except IOError:
						log.error("Could not write to lame!")
						self.finished = True
						break
				else:
					try:
						tmp = 0
						for chunk in data.render(self.stream_chunk_size):
							try:
								samples = len(chunk)
								self.buffered += samples
								self.lame_input_length += samples
								tmp += samples
								chunk.tofile(self.lame.stdin)
							except IOError:
								log.error("Could not write to lame!")
								self.finished = True
								break
						self.delta += tmp - data.samples
						log.debug("Current delta: %d samples.", self.delta)
						#   Note: this delta will cause drift of 1 second/month.
						#   TODO: Fix it. Eventually.
					except Exception:
						log.error("Couldn't render segment due to:\n%s",
								traceback.format_exc())
			except Exception:
				log.critical("Failed to write to Lame:\n%s",
							 traceback.format_exc())
			finally:
				self.encode.release()
		log.critical("Encoder finishing!")

	#   TODO: Extend me to work for all samplerates
	def start(self, *args, **kwargs):
		call = ["lame"]
		call.append('-r')
		if self.input_wordlength != 16:
			call.extend(["--bitwidth", str(self.input_wordlength)])
		call.extend(self.preset.split())
		call.extend(["-", "-"])
		self.lame = subprocess.Popen(
			call,
			stdin=subprocess.PIPE,
			stdout=subprocess.PIPE,
			stderr=subprocess.PIPE
		)
		threading.Thread.start(self, *args, **kwargs)

	def ensure_is_alive(self):
		if self.finished:
			return False
		if self.is_alive():
			return True
		try:
			self.start()
			return True
		except Exception:
			return False

	def run(self, *args, **kwargs):
		try:
			last = None
			lag = 0
			while True:
				timing = float(SAMPLES_PER_FRAME) / self.samplerate

				header = self.lame.stdout.read(HEADER_SIZE)
				if len(header) == HEADER_SIZE:
					frame_len = frame_length(header) - HEADER_SIZE
					frame = self.lame.stdout.read(frame_len)
					buf = header + frame
					if len(frame) == frame_len:
						self.buffered -= SAMPLES_PER_FRAME
				else:
					buf = header

				if self.buffered < (self.safety_buffer * self.samplerate):
					self.ready.release()
				if len(buf):
					self.out_samples += SAMPLES_PER_FRAME
					if self.ofile:
						self.ofile.write(buf)
						self.ofile.flush()
					if self.callback:
						self.callback(False)
					if self.markers and self.markers[0][0] <= self.out_samples:
						msg = self.markers.pop(0)[1];
						log.info("Adding marker: %r", msg)
						if isinstance(msg, unicode): msg=msg.encode("utf-8")
						msg = "\3" + msg # Prepend the UTF-8 encoding marker. If msg was already bytes, it MUST be valid UTF-8.
						frame = "TIT2"+synchsafe(len(msg))+"\0\0"+msg;
						self.oqueue.put("ID3\3\0\0"+synchsafe(len(frame))+frame)
					if self.oqueue:
						self.oqueue.put(buf)
					if self.real_time and self.sent:
						now = time.time()
						if last:
							delta = (now - last - timing)
							lag += delta
							if lag < timing:
								time.sleep(max(0, timing - delta))
						last = now
					self.sent = True
				else:
					if self.callback:
						self.callback(True)
					break
			self.lame.wait()
		except Exception:
			log.error(traceback.format_exc())
			self.finish()
			raise

	def finish(self):
		"""
			Closes input stream to LAME and waits for the last frame(s) to
			finish encoding. Returns LAME's return value code.
		"""
		if self.lame:
			self.__write_queue.put(None)
			self.encode.acquire()
			self.lame.stdin.close()
			self.join()
			self.finished = True
			return self.lame.returncode
		return -1

if __name__ == "__main__":
	import wave
	f = wave.open("test.wav")
	a = numpy.frombuffer(f.readframes(f.getnframes()),
						 dtype=numpy.int16).reshape((-1, 2))

	s = time.time()
	print "Encoding test.wav to testout.mp3..."
	encoder = Lame(ofile=open('testout.mp3', 'w'))
	encoder.safety_buffer = 30
	encoder.start()
	encoder.add_pcm(a)
	encoder.finish()
	s = time.time() - s
	print "Took %2.2fs" % s
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright (c) 2012-2013 SEOmoz, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

'''This is a module for dealing with urls. In particular, sanitizing them.'''

import re
import codecs
import urllib
try:
    import urlparse
except ImportError:  # pragma: no cover
    # Python 3 support
    import urllib.parse as urlparse

# For publicsuffix utilities
from publicsuffix import PublicSuffixList
psl = PublicSuffixList()

# Come codes that we'll need
IDNA = codecs.lookup('idna')
UTF8 = codecs.lookup('utf-8')
ASCII = codecs.lookup('ascii')
W1252 = codecs.lookup('windows-1252')

# The default ports associated with each scheme
PORTS = {
    'http': 80,
    'https': 443
}

def parse(url, encoding='utf-8'):
    '''Parse the provided url string and return an URL object'''
    return URL.parse(url, encoding)

class URL(object):
    '''
    For more information on how and what we parse / sanitize:
        http://tools.ietf.org/html/rfc1808.html
    The more up-to-date RFC is this one:
        http://www.ietf.org/rfc/rfc3986.txt
    '''

    # Via http://www.ietf.org/rfc/rfc3986.txt
    GEN_DELIMS = ":/?#[]@"
    SUB_DELIMS = "!$&'()*+,;="
    ALPHA = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    DIGIT = "0123456789"
    UNRESERVED = ALPHA + DIGIT + "-._~"
    RESERVED = GEN_DELIMS + SUB_DELIMS
    PCHAR = UNRESERVED + SUB_DELIMS + ":@"
    PATH = PCHAR + "/"
    QUERY = PCHAR + "/?"
    FRAGMENT = PCHAR + "/?"
    USERINFO = UNRESERVED + SUB_DELIMS + ":"

    PERCENT_ESCAPING_RE = re.compile('(%([a-fA-F0-9]{2})|.)', re.S)

    @classmethod
    def parse(cls, url, encoding):
        '''Parse the provided url, and return a URL instance'''
        if isinstance(url, str):
            parsed = urlparse.urlparse(
                url.decode(encoding).encode('utf-8'))
        else:
            parsed = urlparse.urlparse(url.encode('utf-8'))

        try:
            port = parsed.port
        except ValueError:
            port = None

        userinfo = parsed.username
        if userinfo and parsed.password:
            userinfo += ':%s' % parsed.password

        return cls(parsed.scheme, parsed.hostname, port,
            parsed.path, parsed.params, parsed.query, parsed.fragment, userinfo)

    def __init__(self, scheme, host, port, path, params, query, fragment, userinfo=None):
        self.scheme = scheme
        self.host = host
        self.port = port
        self.path = path or '/'
        self.params = re.sub(r'^;+', '', str(params))
        self.params = re.sub(r'^;|;$', '', re.sub(r';{2,}', ';', self.params))
        # Strip off extra leading ?'s
        self.query = re.sub(r'^\?+', '', str(query))
        self.query = re.sub(r'^&|&$', '', re.sub(r'&{2,}', '&', self.query))
        self.fragment = fragment
        self.userinfo = userinfo

    def copy(self):
        '''Return a new instance of an identical URL.'''
        return URL(
            self.scheme,
            self.host,
            self.port,
            self.path,
            self.params,
            self.query,
            self.fragment,
            self.userinfo)

    def equiv(self, other):
        '''Return true if this url is equivalent to another'''
        if isinstance(other, basestring):
            _other = self.parse(other, 'utf-8')
        else:
            _other = self.parse(other.utf8, 'utf-8')

        _self = self.parse(self.utf8, 'utf-8')
        _self.canonical().defrag().abspath().escape().punycode()
        _other.canonical().defrag().abspath().escape().punycode()

        result = (
            _self.scheme    == _other.scheme    and
            _self.host      == _other.host      and
            _self.path      == _other.path      and
            _self.params    == _other.params    and
            _self.query     == _other.query)

        if result:
            if _self.port and not _other.port:
                # Make sure _self.port is the default for the scheme
                return _self.port == PORTS.get(_self.scheme, None)
            elif _other.port and not _self.port:
                # Make sure _other.port is the default for the scheme
                return _other.port == PORTS.get(_other.scheme, None)
            else:
                return _self.port == _other.port
        else:
            return False

    def __eq__(self, other):
        '''Return true if this url is /exactly/ equal to another'''
        if isinstance(other, basestring):
            return self.__eq__(self.parse(other, 'utf-8'))
        return (
            self.scheme    == other.scheme    and
            self.host      == other.host      and
            self.path      == other.path      and
            self.port      == other.port      and
            self.params    == other.params    and
            self.query     == other.query     and
            self.fragment  == other.fragment  and
            self.userinfo  == other.userinfo)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __str__(self):
        return self.utf8

    def __repr__(self):
        return '<url.URL object "%s" >' % self.utf8

    def canonical(self):
        '''Canonicalize this url. This includes reordering parameters and args
        to have a consistent ordering'''
        self.query = '&'.join(sorted([q for q in self.query.split('&')]))
        self.params = ';'.join(sorted([q for q in self.params.split(';')]))
        return self

    def defrag(self):
        '''Remove the fragment from this url'''
        self.fragment = None
        return self

    def deparam(self, params):
        '''Strip any of the provided parameters out of the url'''
        lowered = set([p.lower() for p in params])
        def function(name, _):
            return name.lower() in lowered
        return self.filter_params(function)

    def filter_params(self, function):
        '''Remove parameters if function(name, value)'''
        def keep(query):
            name, _, value = query.partition('=')
            return not function(name, value)
        self.query = '&'.join(q for q in self.query.split('&') if q and keep(q))
        self.params = ';'.join(q for q in self.params.split(';') if q and keep(q))
        return self

    def deuserinfo(self):
        '''Remove any userinfo'''
        self.userinfo = None
        return self

    def abspath(self):
        '''Clear out any '..' and excessive slashes from the path'''
        # Remove double forward-slashes from the path
        path = re.sub(r'\/{2,}', '/', self.path)
        # With that done, go through and remove all the relative references
        unsplit = []
        directory = False
        for part in path.split('/'):
            # If we encounter the parent directory, and there's
            # a segment to pop off, then we should pop it off.
            if part == '..' and (not unsplit or unsplit.pop() != None):
                directory = True
            elif part != '.':
                unsplit.append(part)
                directory = False
            else:
                directory = True

        # With all these pieces, assemble!
        if directory:
            # If the path ends with a period, then it refers to a directory,
            # not a file path
            self.path = '/'.join(unsplit) + '/'
        else:
            self.path = '/'.join(unsplit)
        return self

    def sanitize(self):
        '''A shortcut to abspath and escape'''
        return self.abspath().escape()

    def remove_default_port(self):
        '''If a port is provided an is the default, remove it.'''
        if self.port and self.scheme and (self.port == PORTS[self.scheme]):
            self.port = None
        return self

    @staticmethod
    def percent_encode(raw, safe):
        if isinstance(raw, unicode):
            raw = UTF8.encode(raw)[0]

        def replacement(match):
            string = match.group(1)
            if len(string) == 1:
                if string in safe:
                    return string
                else:
                    return '%%%02X' % ord(string)
            else:
                # Replace any escaped entities with their equivalent if needed.
                character = chr(int(match.group(2), 16))
                if (character in safe) and not (character in URL.RESERVED):
                    return character
                return string.upper()

        return URL.PERCENT_ESCAPING_RE.sub(replacement, raw)

    def escape(self, strict=False):
        '''Make sure that the path is correctly escaped'''
        if strict:
            self.path = self.percent_encode(self.path, URL.PATH)
            self.query = self.percent_encode(self.query, URL.QUERY)
            self.params = self.percent_encode(self.params, URL.QUERY)
            if self.userinfo:
                self.userinfo = self.percent_encode(self.userinfo, URL.USERINFO)
            return self
        else:
            self.path = urllib.quote(
                urllib.unquote(self.path), safe=URL.PATH)
            # Safe characters taken from:
            #    http://tools.ietf.org/html/rfc3986#page-50
            self.query = urllib.quote(urllib.unquote(self.query),
                safe=URL.QUERY)
            # The safe characters for URL parameters seemed a little more vague.
            # They are interpreted here as *pchar despite this page, since the
            # updated RFC seems to offer no replacement
            #    http://tools.ietf.org/html/rfc3986#page-54
            self.params = urllib.quote(urllib.unquote(self.params),
                safe=URL.QUERY)
            if self.userinfo:
                self.userinfo = urllib.quote(urllib.unquote(self.userinfo),
                    safe=URL.USERINFO)
            return self

    def unescape(self):
        '''Unescape the path'''
        self.path = urllib.unquote(self.path)
        return self

    def encode(self, encoding):
        '''Return the url in an arbitrary encoding'''
        netloc = self.host or ''
        if self.port:
            netloc += (':' + str(self.port))

        if self.userinfo is not None:
            netloc = '%s@%s' % (self.userinfo, netloc)

        result = urlparse.urlunparse((str(self.scheme), str(netloc),
            str(self.path), str(self.params), str(self.query),
            self.fragment))
        return result.decode('utf-8').encode(encoding)

    def relative(self, path, encoding='utf-8', errors='replace'):
        '''Evaluate the new path relative to the current url'''
        if not isinstance(path, str):
            newurl = urlparse.urljoin(self.utf8, path.encode('utf-8', errors))
        else:
            newurl = urlparse.urljoin(
                self.utf8, path.decode(encoding, errors).encode('utf-8', errors))
        return URL.parse(newurl, 'utf-8')

    def punycode(self):
        '''Convert to punycode hostname'''
        if self.host:
            self.host = IDNA.encode(self.host.decode('utf-8'))[0]
            return self
        raise TypeError('Cannot punycode a relative url (%s)' % repr(self))

    def unpunycode(self):
        '''Convert to an unpunycoded hostname'''
        if self.host:
            self.host = IDNA.decode(
                self.host.decode('utf-8'))[0].encode('utf-8')
            return self
        raise TypeError('Cannot unpunycode a relative url (%s)' % repr(self))

    ###########################################################################
    # Information about the domain
    ###########################################################################
    @property
    def hostname(self):
        '''Return the hostname of the url.'''
        return self.host or ''

    @property
    def pld(self):
        '''Return the 'pay-level domain' of the url
            (http://moz.com/blog/what-the-heck-should-we-call-domaincom)'''
        if self.host:
            return psl.get_public_suffix(self.host)
        return ''

    @property
    def tld(self):
        '''Return the top-level domain of a url'''
        if self.host:
            return '.'.join(self.pld.split('.')[1:])
        return ''

    ###########################################################################
    # Information about the type of url it is
    ###########################################################################
    @property
    def absolute(self):
        '''Return True if this is a fully-qualified URL with a hostname and
        everything'''
        return bool(self.host)

    ###########################################################################
    # Get a string representation. These methods can't be chained, as they
    # return strings
    ###########################################################################
    @property
    def unicode(self):
        '''Return a unicode version of this url'''
        return self.encode('utf-8').decode('utf-8')

    @property
    def utf8(self):
        '''Return a utf-8 version of this url'''
        return self.encode('utf-8')
<commit_msg>Fixed psl implementation.<commit_after>#!/usr/bin/env python
#
# Copyright (c) 2012-2013 SEOmoz, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

'''This is a module for dealing with urls. In particular, sanitizing them.'''

import re
import codecs
import urllib
try:
    import urlparse
except ImportError:  # pragma: no cover
    # Python 3 support
    import urllib.parse as urlparse

# For publicsuffix utilities
from publicsuffix import PublicSuffixList
psl = PublicSuffixList()

# Come codes that we'll need
IDNA = codecs.lookup('idna')
UTF8 = codecs.lookup('utf-8')
ASCII = codecs.lookup('ascii')
W1252 = codecs.lookup('windows-1252')

# The default ports associated with each scheme
PORTS = {
    'http': 80,
    'https': 443
}

def parse(url, encoding='utf-8'):
    '''Parse the provided url string and return an URL object'''
    return URL.parse(url, encoding)

class URL(object):
    '''
    For more information on how and what we parse / sanitize:
        http://tools.ietf.org/html/rfc1808.html
    The more up-to-date RFC is this one:
        http://www.ietf.org/rfc/rfc3986.txt
    '''

    # Via http://www.ietf.org/rfc/rfc3986.txt
    GEN_DELIMS = ":/?#[]@"
    SUB_DELIMS = "!$&'()*+,;="
    ALPHA = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
    DIGIT = "0123456789"
    UNRESERVED = ALPHA + DIGIT + "-._~"
    RESERVED = GEN_DELIMS + SUB_DELIMS
    PCHAR = UNRESERVED + SUB_DELIMS + ":@"
    PATH = PCHAR + "/"
    QUERY = PCHAR + "/?"
    FRAGMENT = PCHAR + "/?"
    USERINFO = UNRESERVED + SUB_DELIMS + ":"

    PERCENT_ESCAPING_RE = re.compile('(%([a-fA-F0-9]{2})|.)', re.S)

    @classmethod
    def parse(cls, url, encoding):
        '''Parse the provided url, and return a URL instance'''
        if isinstance(url, str):
            parsed = urlparse.urlparse(
                url.decode(encoding).encode('utf-8'))
        else:
            parsed = urlparse.urlparse(url.encode('utf-8'))

        try:
            port = parsed.port
        except ValueError:
            port = None

        userinfo = parsed.username
        if userinfo and parsed.password:
            userinfo += ':%s' % parsed.password

        return cls(parsed.scheme, parsed.hostname, port,
            parsed.path, parsed.params, parsed.query, parsed.fragment, userinfo)

    def __init__(self, scheme, host, port, path, params, query, fragment, userinfo=None):
        self.scheme = scheme
        self.host = host
        self.port = port
        self.path = path or '/'
        self.params = re.sub(r'^;+', '', str(params))
        self.params = re.sub(r'^;|;$', '', re.sub(r';{2,}', ';', self.params))
        # Strip off extra leading ?'s
        self.query = re.sub(r'^\?+', '', str(query))
        self.query = re.sub(r'^&|&$', '', re.sub(r'&{2,}', '&', self.query))
        self.fragment = fragment
        self.userinfo = userinfo

    def copy(self):
        '''Return a new instance of an identical URL.'''
        return URL(
            self.scheme,
            self.host,
            self.port,
            self.path,
            self.params,
            self.query,
            self.fragment,
            self.userinfo)

    def equiv(self, other):
        '''Return true if this url is equivalent to another'''
        if isinstance(other, basestring):
            _other = self.parse(other, 'utf-8')
        else:
            _other = self.parse(other.utf8, 'utf-8')

        _self = self.parse(self.utf8, 'utf-8')
        _self.canonical().defrag().abspath().escape().punycode()
        _other.canonical().defrag().abspath().escape().punycode()

        result = (
            _self.scheme    == _other.scheme    and
            _self.host      == _other.host      and
            _self.path      == _other.path      and
            _self.params    == _other.params    and
            _self.query     == _other.query)

        if result:
            if _self.port and not _other.port:
                # Make sure _self.port is the default for the scheme
                return _self.port == PORTS.get(_self.scheme, None)
            elif _other.port and not _self.port:
                # Make sure _other.port is the default for the scheme
                return _other.port == PORTS.get(_other.scheme, None)
            else:
                return _self.port == _other.port
        else:
            return False

    def __eq__(self, other):
        '''Return true if this url is /exactly/ equal to another'''
        if isinstance(other, basestring):
            return self.__eq__(self.parse(other, 'utf-8'))
        return (
            self.scheme    == other.scheme    and
            self.host      == other.host      and
            self.path      == other.path      and
            self.port      == other.port      and
            self.params    == other.params    and
            self.query     == other.query     and
            self.fragment  == other.fragment  and
            self.userinfo  == other.userinfo)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __str__(self):
        return self.utf8

    def __repr__(self):
        return '<url.URL object "%s" >' % self.utf8

    def canonical(self):
        '''Canonicalize this url. This includes reordering parameters and args
        to have a consistent ordering'''
        self.query = '&'.join(sorted([q for q in self.query.split('&')]))
        self.params = ';'.join(sorted([q for q in self.params.split(';')]))
        return self

    def defrag(self):
        '''Remove the fragment from this url'''
        self.fragment = None
        return self

    def deparam(self, params):
        '''Strip any of the provided parameters out of the url'''
        lowered = set([p.lower() for p in params])
        def function(name, _):
            return name.lower() in lowered
        return self.filter_params(function)

    def filter_params(self, function):
        '''Remove parameters if function(name, value)'''
        def keep(query):
            name, _, value = query.partition('=')
            return not function(name, value)
        self.query = '&'.join(q for q in self.query.split('&') if q and keep(q))
        self.params = ';'.join(q for q in self.params.split(';') if q and keep(q))
        return self

    def deuserinfo(self):
        '''Remove any userinfo'''
        self.userinfo = None
        return self

    def abspath(self):
        '''Clear out any '..' and excessive slashes from the path'''
        # Remove double forward-slashes from the path
        path = re.sub(r'\/{2,}', '/', self.path)
        # With that done, go through and remove all the relative references
        unsplit = []
        directory = False
        for part in path.split('/'):
            # If we encounter the parent directory, and there's
            # a segment to pop off, then we should pop it off.
            if part == '..' and (not unsplit or unsplit.pop() != None):
                directory = True
            elif part != '.':
                unsplit.append(part)
                directory = False
            else:
                directory = True

        # With all these pieces, assemble!
        if directory:
            # If the path ends with a period, then it refers to a directory,
            # not a file path
            self.path = '/'.join(unsplit) + '/'
        else:
            self.path = '/'.join(unsplit)
        return self

    def sanitize(self):
        '''A shortcut to abspath and escape'''
        return self.abspath().escape()

    def remove_default_port(self):
        '''If a port is provided an is the default, remove it.'''
        if self.port and self.scheme and (self.port == PORTS[self.scheme]):
            self.port = None
        return self

    @staticmethod
    def percent_encode(raw, safe):
        if isinstance(raw, unicode):
            raw = UTF8.encode(raw)[0]

        def replacement(match):
            string = match.group(1)
            if len(string) == 1:
                if string in safe:
                    return string
                else:
                    return '%%%02X' % ord(string)
            else:
                # Replace any escaped entities with their equivalent if needed.
                character = chr(int(match.group(2), 16))
                if (character in safe) and not (character in URL.RESERVED):
                    return character
                return string.upper()

        return URL.PERCENT_ESCAPING_RE.sub(replacement, raw)

    def escape(self, strict=False):
        '''Make sure that the path is correctly escaped'''
        if strict:
            self.path = self.percent_encode(self.path, URL.PATH)
            self.query = self.percent_encode(self.query, URL.QUERY)
            self.params = self.percent_encode(self.params, URL.QUERY)
            if self.userinfo:
                self.userinfo = self.percent_encode(self.userinfo, URL.USERINFO)
            return self
        else:
            self.path = urllib.quote(
                urllib.unquote(self.path), safe=URL.PATH)
            # Safe characters taken from:
            #    http://tools.ietf.org/html/rfc3986#page-50
            self.query = urllib.quote(urllib.unquote(self.query),
                safe=URL.QUERY)
            # The safe characters for URL parameters seemed a little more vague.
            # They are interpreted here as *pchar despite this page, since the
            # updated RFC seems to offer no replacement
            #    http://tools.ietf.org/html/rfc3986#page-54
            self.params = urllib.quote(urllib.unquote(self.params),
                safe=URL.QUERY)
            if self.userinfo:
                self.userinfo = urllib.quote(urllib.unquote(self.userinfo),
                    safe=URL.USERINFO)
            return self

    def unescape(self):
        '''Unescape the path'''
        self.path = urllib.unquote(self.path)
        return self

    def encode(self, encoding):
        '''Return the url in an arbitrary encoding'''
        netloc = self.host or ''
        if self.port:
            netloc += (':' + str(self.port))

        if self.userinfo is not None:
            netloc = '%s@%s' % (self.userinfo, netloc)

        result = urlparse.urlunparse((str(self.scheme), str(netloc),
            str(self.path), str(self.params), str(self.query),
            self.fragment))
        return result.decode('utf-8').encode(encoding)

    def relative(self, path, encoding='utf-8', errors='replace'):
        '''Evaluate the new path relative to the current url'''
        if not isinstance(path, str):
            newurl = urlparse.urljoin(self.utf8, path.encode('utf-8', errors))
        else:
            newurl = urlparse.urljoin(
                self.utf8, path.decode(encoding, errors).encode('utf-8', errors))
        return URL.parse(newurl, 'utf-8')

    def punycode(self):
        '''Convert to punycode hostname'''
        if self.host:
            self.host = IDNA.encode(self.host.decode('utf-8'))[0]
            return self
        raise TypeError('Cannot punycode a relative url (%s)' % repr(self))

    def unpunycode(self):
        '''Convert to an unpunycoded hostname'''
        if self.host:
            self.host = IDNA.decode(
                self.host.decode('utf-8'))[0].encode('utf-8')
            return self
        raise TypeError('Cannot unpunycode a relative url (%s)' % repr(self))

    ###########################################################################
    # Information about the domain
    ###########################################################################
    @property
    def hostname(self):
        '''Return the hostname of the url.'''
        return self.host or ''

    @property
    def pld(self):
        '''Return the 'pay-level domain' of the url
            (http://moz.com/blog/what-the-heck-should-we-call-domaincom)'''
        if self.host:
            if not 'xn--' in self.host:
                return psl.get_public_suffix(self.host)
            unpunyhost = IDNA.decode(self.host)[0]
            tld = psl.get_public_suffix(unpunyhost)
            return IDNA.encode(tld)[0]

        return ''

    @property
    def tld(self):
        '''Return the top-level domain of a url'''
        if self.host:
            return '.'.join(self.pld.split('.')[1:])
        return ''

    ###########################################################################
    # Information about the type of url it is
    ###########################################################################
    @property
    def absolute(self):
        '''Return True if this is a fully-qualified URL with a hostname and
        everything'''
        return bool(self.host)

    ###########################################################################
    # Get a string representation. These methods can't be chained, as they
    # return strings
    ###########################################################################
    @property
    def unicode(self):
        '''Return a unicode version of this url'''
        return self.encode('utf-8').decode('utf-8')

    @property
    def utf8(self):
        '''Return a utf-8 version of this url'''
        return self.encode('utf-8')
<|endoftext|>


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

import sys, os, json, platform
import subprocess, string, re

_json_config_file = os.getenv('USE_CONFIG_FILE')
_targets_folder = os.getenv('USE_TARGETS_FOLDER')
_targets = {}
_rcfile = ""
_arguments = sys.argv[1:]
_configure = False
_switches = []
_ask_for_ssh_keys = False

POSSIBLE_SWITCHES = ['--keep', '--config', '--configure', '--edit', '--conf', '--help', '-h', '--bash-autocomplete-helper']

if not _json_config_file:
    print "Configuration file not found!\nSet the env variable USE_CONFIG_FILE, point it to your json file.\n"
    sys.exit(-1)

if not _targets_folder:
    print "Use folder not found!\nSet env variable USE_TARGETS_FOLDER, point it to your folder with env scripts.\n"
    sys.exit(-1)

def osType(): # returns 'nt' or 'posix'
    return os.name

def platformName(): # returns 'Windows', 'Linux' or 'Darwin'
    return platform.system()

def isWindows():
    return platform.system() == "Windows"

def isLinux():
    return platform.system() == "Linux"

def isBash():
    return "bash" in os.getenv("SHELL")

def list_separator():
    if os.name == 'nt':
        return ';'
    return ':'

def to_native_path(path):
    path = os.path.abspath(path)
    if path.endswith("\\") or (path.endswith('/') and path != '/'):
        path = path[:-1]
    return path

def fill_placeholders(value):
    placeholders = re.findall('\$\{(.*?)\}', value) # searches for ${foo}
    for placeholder in placeholders:
        if placeholder in os.environ:
            value = value.replace("${" + placeholder + "}", os.environ[placeholder])

    return value

class EnvVariable:
    def __init__(self):
        self.name = ""
        self.value = ""
        self.values = []
    def isPath(self):
        return self.values or '/' in self.value or '\\' in self.value

class Target:
    def __init__(self, tname):
        self.name = tname
        self.uses = []
        self.uses_after = []
        self.cwd = ""
        self.hidden = False
        self.yakuake_tab_name = ""
        self.platforms = []
        self.variables = []
        self.arg = ""

        self.loadJson()

    def jsonFileName(self):
        return _targets_folder + "/../unified/" + self.name + ".json"

    def env_var_from_json(self, json):
        var = EnvVariable()
        key = json.keys()[0]
        var.name = key
        value = json[key]
        value_is_list = type(value) == type([])

        if value_is_list:
            var.values = value
        else:
            var.value = str(value)

        return var

    def loadJson(self):
        if not os.path.exists(self.jsonFileName()):
            return False
        return self.loadJsonFile(self.jsonFileName())

    def isGeneric(self):
        return "%" in self.name

    def displayName(self):
        if self.isGeneric() and self.arg:
            return self.name.replace("%", self.arg)
        return self.name

    def loadJsonFile(self, filename):
        if not os.path.exists(filename):
            print "File doesn't exist: " + filename
            return False

        f = open(filename, 'r')
        # print "Processing " + filename
        contents = f.read()
        f.close()
        decoded = json.loads(contents)

        # first source 'nt' and 'posix'
        if osType() in decoded:
            for env_var in decoded[osType()]:
                self.variables.append(self.env_var_from_json(env_var))

        # now source 'Linux', 'Darwin'or 'Windows', which have precedence
        if platformName() in decoded:
            for env_var in decoded[platformName()]:
                self.variables.append(self.env_var_from_json(env_var))

        # Source the platform-independent variables
        if "any" in decoded:
            for env_var in decoded["any"]:
                self.variables.append(self.env_var_from_json(env_var))

        if "includes" in decoded:
            for include in decoded['includes']:
                if not self.loadJsonFile(fill_placeholders(include)):
                    return False
        return True

def printUsage():
    print "Usage:"
    print sys.argv[0] + " <target>\n"

    print "Available targets:\n"
    for target in _targets:
        t = _targets[target]
        if not t.hidden:
            print "  " + target

    sys.exit(1)

def cleanup_cwd(cwd):
    if cwd.startswith("$"):
        env_var_name = cwd[1:]
        if env_var_name and env_var_name in os.environ:
            return os.environ[env_var_name]
    return cwd

def loadJson():
    f = open(_json_config_file, 'r')
    contents = f.read()
    f.close()

    decoded = json.loads(contents)

    global _rcfile, _ask_for_ssh_keys

    if "targets" in decoded:
        for target in decoded['targets']:
            name = ""
            if "name" in target:
                name = target['name']
            else:
                print "Missing name for target"
                return False

            t = Target(name)

            if "uses" in target:
                t.uses = target['uses']

            if "uses_after" in target:
                t.uses_after = target['uses_after']

            if "cwd" in target:
                t.cwd = target['cwd']

            if "rename_yakuake_to" in target:
                t.yakuake_tab_name = target['rename_yakuake_to']

            if "hidden" in target:
                t.hidden = target['hidden']

            _targets[t.name] = t

    if "rcfile" in decoded:
        _rcfile = decoded['rcfile']
        if not os.path.exists(_rcfile):
            print "Requested rcfile doesn't exist: " + _rcfile
            return False

    if "ask_for_ssh_keys" in decoded:
        _ask_for_ssh_keys = decoded['ask_for_ssh_keys']

    return True;

def getGenericTargetAndArg(name):
    tokens = re.findall('(.*)-(.*)', name)
    tokens = tokens[0]

    if len(tokens) == 2:
        return { "name" : tokens[0], "arg" : tokens[1] }

    return {}

def getTarget(name):
    if name in _targets:
        return _targets[name]

    genericTarget = getGenericTargetAndArg(name)

    if genericTarget and genericTarget["name"] + "-%" in _targets:
        target = _targets[genericTarget["name"] + "-%"]
        target.arg = genericTarget["arg"]
        return target

    print "Unknown target: " + name
    printUsage()

def source_single_json(target):
    for v in target.variables:
        if not v.name:
            continue

        if v.value:
            if v.value == "USE_ARG":
                value = target.arg
            else:
                value = fill_placeholders(v.value)
            if v.isPath():
                value = to_native_path(value)
            os.environ[v.name] = value
        else: # list case
            value = list_separator()
            for list_token in v.values:
                list_token = fill_placeholders(list_token)
                if v.isPath():
                    list_token = to_native_path(list_token)
                value = value + list_separator() + list_token

            os.environ[v.name] = value.strip(list_separator())

def source_single_file(filename):
    command = ""

    filename_cmd = filename

    shell = shellForOS(filename)
    if shell == 'cmd':
        command = ['cmd', '/C', filename_cmd + ' && set']
        # os.environ['PROMPT'] = ""
    else:
        command = [shell, '-c', 'source ' + filename_cmd + ' && env']

    # print "Sourcing " + filename
    proc = subprocess.Popen(command, stdout = subprocess.PIPE)

    print "Sourcing " + to_native_path(filename_cmd)

    for line in proc.stdout:
        (key, _, value) = line.partition("=")
        if key and not key.startswith('BASH_FUNC_'):
            try:
                os.environ[key] = value.strip()
            except:
                print "Error importing key=" + key + "; with value=" + value.strip()
                throw
    proc.communicate()

    return True

def extensionForScript():
    if isWindows():
        return ".bat"
    return ".source"

def filenameForTarget(target):
    filename = _targets_folder + "/" + target.name + extensionForScript()
    if os.path.exists(target.jsonFileName()) or target.isGeneric():
        if os.path.exists(filename):
            print "Favoring .json over " + filename

        return target.jsonFileName()

    return filename

def currentTargets():
    targets = os.getenv('USE_CURRENT_TARGETS')
    if not targets:
        return []
    return targets.split(';')

def currentTargetsStr():
    return string.join(currentTargets(), ' ');

def shellForOS(filename = ""):
    # .bat files are always sourced by cmd. Use .json if you don't like this
    if filename.endswith(".bat") or filename.endswith(".cmd"):
        return 'cmd'

    envShell = os.getenv('SHELL')

    if envShell and envShell != '/bin/false':
        return envShell

    if isWindows():
        return 'cmd'

    return 'bash'

def run_shell(cwd):
    cmd = ""
    shell =  shellForOS()
    if _rcfile and 'bash' in shell:
        cmd = shell + " --rcfile " + _rcfile
    else:
        cmd = shell

    old_cwd = ""
    if cwd:
        old_cwd = os.getcwd()
        os.chdir(cwd)
    result = True
    try:
        # print 'Running ' + cmd
        result = os.system(cmd) == 0
    except:
        pass

    if cwd:
        os.chdir(old_cwd)

    return result

def is_sourced(targetName):
    if targetName not in currentTargets():
        return False

    return True

def source_target(target):
    if target.name in currentTargets():
        return True

    for targetName in target.uses:
        if not source_target(getTarget(targetName)):
            return False

    filename = filenameForTarget(target)

    if filename.endswith(".json"):
        arg = ""
        if target.arg:
            arg = " " + target.arg
        print "Sourcing " + to_native_path(filename) + arg
        source_single_json(target)
    else:
        if os.path.exists(filename):
            source_single_file(filename)

    newCurTargets = string.join(currentTargets(), ';')

    os.environ['USE_CURRENT_TARGETS'] = newCurTargets + ";" + target.displayName()

    for targetName in target.uses_after:
        if not source_target(getTarget(targetName)):
            return False

    return True

def reset_env():
    os.environ['USE_CURRENT_TARGETS'] = ""
    return source_target(getTarget("default"))

def use_target(target):
    global _switches
    if is_sourced(target.name):
        return True

    if '--keep' not in _switches:
        if not reset_env(): # source default.source
            return False

    # run qdbus before sourcing, otherwise it might use an incompatible Qt
    must_restore_yakuake = False
    if target.yakuake_tab_name and isLinux():
        os.system("rename_yatab.sh " + target.yakuake_tab_name)
        must_restore_yakuake = True

    success = False
    if source_target(target):
        success = run_shell(cleanup_cwd(target.cwd)) # this hangs here until user exits bash
    else:
        success = False

    if must_restore_yakuake:
        os.system("rename_yatab.sh Shell")

    return success

def editor():
    ed = os.getenv('USE_EDITOR')
    if ed:
        return ed
    return 'kate'

def open_editor(filename):
    return os.system(editor() + " " + filename) == 0

def process_arguments():
    global _switches
    for a in _arguments:
        if a in POSSIBLE_SWITCHES:
            _arguments.remove(a)
            _switches.append(a)
        elif a.startswith('--') and not '--bash-autocomplete-helper' in _arguments:
            print "Invalid switch: " + a
            sys.exit(-1)

def source_default():
    t = Target("default")
    _targets[t.name] = t

def ask_for_ssh_keys():
    try:
        subprocess.check_output(["ssh-add", "-L"]) == 0
    except:
        return os.system("ssh-add") == 0

    return True

process_arguments()

if '--config' in _switches or '--configure' in _switches or '--conf' in _switches:
    open_editor(_json_config_file)
    sys.exit(1)

source_default()

if not loadJson():
    print "Error loading json"
    sys.exit(1)

if len(sys.argv) == 1:
    print currentTargetsStr()
    sys.exit(-1)

_targetName = sys.argv[1]

if '--edit' in _switches:
    filename = filenameForTarget(Target(_targetName))
    print "Opening editor for " + filename
    if not open_editor(filename):
        print "Error opening editor"
    sys.exit(0)

if '--help' in _switches or '-h' in _switches:
    printUsage()
    sys.exit(0)

if '--bash-autocomplete-helper' in _switches:
    result = ''
    wordBeginning = ""
    if _arguments:
        wordBeginning = _arguments[0]

    for target in _targets:
        t = _targets[target]
        if not t.hidden and target.startswith(wordBeginning):
            result = result + ' ' + target

    if wordBeginning.startswith('-'):
        # Let's only spam the completion with switches if the user already typed -
        for s in POSSIBLE_SWITCHES:
            if s.startswith(wordBeginning):
                result = result + ' ' + s

    print result.strip()
    sys.exit(0)

if _ask_for_ssh_keys:
    ask_for_ssh_keys()

use_target(getTarget(_targetName))
<commit_msg>Fix not allowing to source the same target twice<commit_after>#!/usr/bin/env python2

import sys, os, json, platform
import subprocess, string, re

_json_config_file = os.getenv('USE_CONFIG_FILE')
_targets_folder = os.getenv('USE_TARGETS_FOLDER')
_targets = {}
_rcfile = ""
_arguments = sys.argv[1:]
_configure = False
_switches = []
_ask_for_ssh_keys = False

POSSIBLE_SWITCHES = ['--keep', '--config', '--configure', '--edit', '--conf', '--help', '-h', '--bash-autocomplete-helper']

if not _json_config_file:
    print "Configuration file not found!\nSet the env variable USE_CONFIG_FILE, point it to your json file.\n"
    sys.exit(-1)

if not _targets_folder:
    print "Use folder not found!\nSet env variable USE_TARGETS_FOLDER, point it to your folder with env scripts.\n"
    sys.exit(-1)

def osType(): # returns 'nt' or 'posix'
    return os.name

def platformName(): # returns 'Windows', 'Linux' or 'Darwin'
    return platform.system()

def isWindows():
    return platform.system() == "Windows"

def isLinux():
    return platform.system() == "Linux"

def isBash():
    return "bash" in os.getenv("SHELL")

def list_separator():
    if os.name == 'nt':
        return ';'
    return ':'

def to_native_path(path):
    path = os.path.abspath(path)
    if path.endswith("\\") or (path.endswith('/') and path != '/'):
        path = path[:-1]
    return path

def fill_placeholders(value):
    placeholders = re.findall('\$\{(.*?)\}', value) # searches for ${foo}
    for placeholder in placeholders:
        if placeholder in os.environ:
            value = value.replace("${" + placeholder + "}", os.environ[placeholder])

    return value

class EnvVariable:
    def __init__(self):
        self.name = ""
        self.value = ""
        self.values = []
    def isPath(self):
        return self.values or '/' in self.value or '\\' in self.value

class Target:
    def __init__(self, tname):
        self.name = tname
        self.uses = []
        self.uses_after = []
        self.cwd = ""
        self.hidden = False
        self.yakuake_tab_name = ""
        self.platforms = []
        self.variables = []
        self.arg = ""

        self.loadJson()

    def jsonFileName(self):
        return _targets_folder + "/../unified/" + self.name + ".json"

    def env_var_from_json(self, json):
        var = EnvVariable()
        key = json.keys()[0]
        var.name = key
        value = json[key]
        value_is_list = type(value) == type([])

        if value_is_list:
            var.values = value
        else:
            var.value = str(value)

        return var

    def loadJson(self):
        if not os.path.exists(self.jsonFileName()):
            return False
        return self.loadJsonFile(self.jsonFileName())

    def isGeneric(self):
        return "%" in self.name

    def displayName(self):
        if self.isGeneric() and self.arg:
            return self.name.replace("%", self.arg)
        return self.name

    def loadJsonFile(self, filename):
        if not os.path.exists(filename):
            print "File doesn't exist: " + filename
            return False

        f = open(filename, 'r')
        # print "Processing " + filename
        contents = f.read()
        f.close()
        decoded = json.loads(contents)

        # first source 'nt' and 'posix'
        if osType() in decoded:
            for env_var in decoded[osType()]:
                self.variables.append(self.env_var_from_json(env_var))

        # now source 'Linux', 'Darwin'or 'Windows', which have precedence
        if platformName() in decoded:
            for env_var in decoded[platformName()]:
                self.variables.append(self.env_var_from_json(env_var))

        # Source the platform-independent variables
        if "any" in decoded:
            for env_var in decoded["any"]:
                self.variables.append(self.env_var_from_json(env_var))

        if "includes" in decoded:
            for include in decoded['includes']:
                if not self.loadJsonFile(fill_placeholders(include)):
                    return False
        return True

def printUsage():
    print "Usage:"
    print sys.argv[0] + " <target>\n"

    print "Available targets:\n"
    for target in _targets:
        t = _targets[target]
        if not t.hidden:
            print "  " + target

    sys.exit(1)

def cleanup_cwd(cwd):
    if cwd.startswith("$"):
        env_var_name = cwd[1:]
        if env_var_name and env_var_name in os.environ:
            return os.environ[env_var_name]
    return cwd

def loadJson():
    f = open(_json_config_file, 'r')
    contents = f.read()
    f.close()

    decoded = json.loads(contents)

    global _rcfile, _ask_for_ssh_keys

    if "targets" in decoded:
        for target in decoded['targets']:
            name = ""
            if "name" in target:
                name = target['name']
            else:
                print "Missing name for target"
                return False

            t = Target(name)

            if "uses" in target:
                t.uses = target['uses']

            if "uses_after" in target:
                t.uses_after = target['uses_after']

            if "cwd" in target:
                t.cwd = target['cwd']

            if "rename_yakuake_to" in target:
                t.yakuake_tab_name = target['rename_yakuake_to']

            if "hidden" in target:
                t.hidden = target['hidden']

            _targets[t.name] = t

    if "rcfile" in decoded:
        _rcfile = decoded['rcfile']
        if not os.path.exists(_rcfile):
            print "Requested rcfile doesn't exist: " + _rcfile
            return False

    if "ask_for_ssh_keys" in decoded:
        _ask_for_ssh_keys = decoded['ask_for_ssh_keys']

    return True;

def getGenericTargetAndArg(name):
    tokens = re.findall('(.*)-(.*)', name)
    tokens = tokens[0]

    if len(tokens) == 2:
        return { "name" : tokens[0], "arg" : tokens[1] }

    return {}

def getTarget(name):
    if name in _targets:
        return _targets[name]

    genericTarget = getGenericTargetAndArg(name)

    if genericTarget and genericTarget["name"] + "-%" in _targets:
        target = _targets[genericTarget["name"] + "-%"]
        target.arg = genericTarget["arg"]
        return target

    print "Unknown target: " + name
    printUsage()

def source_single_json(target):
    for v in target.variables:
        if not v.name:
            continue

        if v.value:
            if v.value == "USE_ARG":
                value = target.arg
            else:
                value = fill_placeholders(v.value)
            if v.isPath():
                value = to_native_path(value)
            os.environ[v.name] = value
        else: # list case
            value = list_separator()
            for list_token in v.values:
                list_token = fill_placeholders(list_token)
                if v.isPath():
                    list_token = to_native_path(list_token)
                value = value + list_separator() + list_token

            os.environ[v.name] = value.strip(list_separator())

def source_single_file(filename):
    command = ""

    filename_cmd = filename

    shell = shellForOS(filename)
    if shell == 'cmd':
        command = ['cmd', '/C', filename_cmd + ' && set']
        # os.environ['PROMPT'] = ""
    else:
        command = [shell, '-c', 'source ' + filename_cmd + ' && env']

    # print "Sourcing " + filename
    proc = subprocess.Popen(command, stdout = subprocess.PIPE)

    print "Sourcing " + to_native_path(filename_cmd)

    for line in proc.stdout:
        (key, _, value) = line.partition("=")
        if key and not key.startswith('BASH_FUNC_'):
            try:
                os.environ[key] = value.strip()
            except:
                print "Error importing key=" + key + "; with value=" + value.strip()
                throw
    proc.communicate()

    return True

def extensionForScript():
    if isWindows():
        return ".bat"
    return ".source"

def filenameForTarget(target):
    filename = _targets_folder + "/" + target.name + extensionForScript()
    if os.path.exists(target.jsonFileName()) or target.isGeneric():
        if os.path.exists(filename):
            print "Favoring .json over " + filename

        return target.jsonFileName()

    return filename

def currentTargets():
    targets = os.getenv('USE_CURRENT_TARGETS')
    if not targets:
        return []
    return targets.split(';')

def currentTargetsStr():
    return string.join(currentTargets(), ' ');

def shellForOS(filename = ""):
    # .bat files are always sourced by cmd. Use .json if you don't like this
    if filename.endswith(".bat") or filename.endswith(".cmd"):
        return 'cmd'

    envShell = os.getenv('SHELL')

    if envShell and envShell != '/bin/false':
        return envShell

    if isWindows():
        return 'cmd'

    return 'bash'

def run_shell(cwd):
    cmd = ""
    shell =  shellForOS()
    if _rcfile and 'bash' in shell:
        cmd = shell + " --rcfile " + _rcfile
    else:
        cmd = shell

    old_cwd = ""
    if cwd:
        old_cwd = os.getcwd()
        os.chdir(cwd)
    result = True
    try:
        # print 'Running ' + cmd
        result = os.system(cmd) == 0
    except:
        pass

    if cwd:
        os.chdir(old_cwd)

    return result

def is_sourced(target):
    if target.displayName() not in currentTargets():
        return False

    return True

def source_target(target):
    if target.name in currentTargets():
        return True

    for targetName in target.uses:
        if not source_target(getTarget(targetName)):
            return False

    filename = filenameForTarget(target)

    if filename.endswith(".json"):
        arg = ""
        if target.arg:
            arg = " " + target.arg
        print "Sourcing " + to_native_path(filename) + arg
        source_single_json(target)
    else:
        if os.path.exists(filename):
            source_single_file(filename)

    newCurTargets = string.join(currentTargets(), ';')

    os.environ['USE_CURRENT_TARGETS'] = newCurTargets + ";" + target.displayName()

    for targetName in target.uses_after:
        if not source_target(getTarget(targetName)):
            return False

    return True

def reset_env():
    os.environ['USE_CURRENT_TARGETS'] = ""
    return source_target(getTarget("default"))

def use_target(target):
    global _switches
    if is_sourced(target):
        return True

    if '--keep' not in _switches:
        if not reset_env(): # source default.source
            return False

    # run qdbus before sourcing, otherwise it might use an incompatible Qt
    must_restore_yakuake = False
    if target.yakuake_tab_name and isLinux():
        os.system("rename_yatab.sh " + target.yakuake_tab_name)
        must_restore_yakuake = True

    success = False
    if source_target(target):
        success = run_shell(cleanup_cwd(target.cwd)) # this hangs here until user exits bash
    else:
        success = False

    if must_restore_yakuake:
        os.system("rename_yatab.sh Shell")

    return success

def editor():
    ed = os.getenv('USE_EDITOR')
    if ed:
        return ed
    return 'kate'

def open_editor(filename):
    return os.system(editor() + " " + filename) == 0

def process_arguments():
    global _switches
    for a in _arguments:
        if a in POSSIBLE_SWITCHES:
            _arguments.remove(a)
            _switches.append(a)
        elif a.startswith('--') and not '--bash-autocomplete-helper' in _arguments:
            print "Invalid switch: " + a
            sys.exit(-1)

def source_default():
    t = Target("default")
    _targets[t.name] = t

def ask_for_ssh_keys():
    try:
        subprocess.check_output(["ssh-add", "-L"]) == 0
    except:
        return os.system("ssh-add") == 0

    return True

process_arguments()

if '--config' in _switches or '--configure' in _switches or '--conf' in _switches:
    open_editor(_json_config_file)
    sys.exit(1)

source_default()

if not loadJson():
    print "Error loading json"
    sys.exit(1)

if len(sys.argv) == 1:
    print currentTargetsStr()
    sys.exit(-1)

_targetName = sys.argv[1]

if '--edit' in _switches:
    filename = filenameForTarget(Target(_targetName))
    print "Opening editor for " + filename
    if not open_editor(filename):
        print "Error opening editor"
    sys.exit(0)

if '--help' in _switches or '-h' in _switches:
    printUsage()
    sys.exit(0)

if '--bash-autocomplete-helper' in _switches:
    result = ''
    wordBeginning = ""
    if _arguments:
        wordBeginning = _arguments[0]

    for target in _targets:
        t = _targets[target]
        if not t.hidden and target.startswith(wordBeginning):
            result = result + ' ' + target

    if wordBeginning.startswith('-'):
        # Let's only spam the completion with switches if the user already typed -
        for s in POSSIBLE_SWITCHES:
            if s.startswith(wordBeginning):
                result = result + ' ' + s

    print result.strip()
    sys.exit(0)

if _ask_for_ssh_keys:
    ask_for_ssh_keys()

use_target(getTarget(_targetName))
<|endoftext|>