<commit_before># -*- coding: utf-8 -*-
# My Model
from utils.ops import scope
import tensorflow as tf
import numpy as np
from  sklearn.datasets import make_blobs

#############################################
#       Deep Adaptive Separator Model       #
#############################################

class KMeans:

	def __init__(self, nb_clusters, centroids_init=None, nb_tries=10, nb_iterations=10, input_tensor=None, latent_space_tensor=None, beta=None):

		self.nb_clusters = nb_clusters
		self.nb_iterations = nb_iterations
		self.nb_tries = nb_tries
		self.latent_space_tensor = latent_space_tensor
		self.beta = beta
		if input_tensor is None:
			self.graph = tf.Graph()
		else:
			self.graph = tf.get_default_graph()

		with self.graph.as_default():
			with tf.name_scope('kmeans'):
				if input_tensor is None:
					# Spectrogram, embeddings
					# shape = [batch, L , E ]
					self.X_in = tf.placeholder("float", [None, None, None], name='Kmeans_input')
				else:
					self.X_in = input_tensor

				# mean, _ = tf.nn.moments(self.X_in, axes=-1, keep_dims=True)
				x_norm = tf.nn.l2_normalize(self.X_in, axis=-1)
				self.b = tf.shape(x_norm)[0]
				self.X = tf.expand_dims(x_norm, 1)
				self.X = tf.tile(self.X, [1, self.nb_tries, 1, 1])

				self.L = tf.shape(self.X)[-2]
				self.E = tf.shape(self.X)[-1]
				self.X = tf.reshape(self.X, [-1, self.L, self.E])

				self.B = tf.shape(self.X)[0]

				self.ones = tf.ones_like(self.X, tf.float32)

				self.shifting = tf.tile(tf.expand_dims(tf.range(self.B)*self.nb_clusters, 1), [1, self.L])

				if centroids_init is None:
					def random_without_replace(b, l):
						a =  np.array([np.random.choice(range(l), size=self.nb_clusters, replace=False) for _ in range(b)])
						return a.astype(np.int32)

					y = tf.py_func(random_without_replace, [self.B, self.L], tf.int32)
					random = tf.reshape(y, [self.B, self.nb_clusters, 1])

					# Take randomly 'nb_clusters' vectors from X
					batch_range = tf.tile(tf.reshape(tf.range(self.B, dtype=tf.int32), shape=[self.B, 1, 1]), [1, self.nb_clusters, 1])
					indices = tf.concat([batch_range, random], axis = 2)
					self.centroid_init = tf.gather_nd(self.X, indices)
				else:
					self.centroids = tf.identity(centroids_init)
					self.centroids = tf.tile(self.centroids, [self.nb_tries, 1 , 1])

				if not self.latent_space_tensor is None:
					log_lst = tf.log(tf.reduce_max(latent_space_tensor, -1, keep_dims=True) / latent_space_tensor)
					self.t = log_lst
					self.notsilent_notry = tf.reshape(tf.cast(log_lst < 2.0, tf.float32), [self.b, self.L, 1])
					self.notsilent = tf.tile(self.notsilent_notry, [self.nb_tries, 1 , 1])
				else:
					self.notsilent = tf.ones([self.B, self.L, 1])

				self.network

	@scope
	def network(self):
		i = tf.constant(0)
		init = [i, self.centroid_init, self.get_labels(self.centroid_init)]
		if self.beta is not None:
			invariant = [i.get_shape(), tf.TensorShape([None, self.nb_clusters, None]), tf.TensorShape([None, None, self.nb_clusters])]
		else:
			invariant = [i.get_shape(), tf.TensorShape([None, self.nb_clusters, None]), tf.TensorShape([None, None])]

		cond = lambda i, c, l: tf.less(i, self.nb_iterations)
		_ , centroids, labels = tf.while_loop(cond, self.body, init, shape_invariants=invariant)
		inertia = self.get_inertia(centroids)

		inertia = tf.reshape(inertia, [self.b, self.nb_tries])
		bests = tf.argmin(inertia, 1, output_type=tf.int32)
		index = bests + tf.range(self.b)*self.nb_tries

		centroids = tf.reshape(centroids, [self.b*self.nb_tries, self.nb_clusters, self.E])
		centroids = tf.gather(centroids, index)

		labels = tf.gather(labels, index)
		return centroids, labels

	def get_inertia(self, centroids):
		with tf.name_scope('inertia'):
			labels = self.get_labels(centroids)
			if self.beta is not None:
				X_ = tf.expand_dims(self.X, 2) # B L 1 E
				centroids_ = tf.expand_dims(centroids, 1) # B 1 C E
				dist = tf.reduce_sum(tf.square(X_ - centroids_), -1) * labels # B C distance to clusters according to soft assignments
				dist = tf.reduce_sum(dist, 1) # Sum over points
				density = tf.reduce_sum(labels, 1) # Compute the soft density of each cluster
				inertia = tf.divide(dist, density)
				inertia = tf.reduce_sum(inertia, -1)
			else:
				# centroids [b*nb_tries, C, E]
				centroids_flattened = tf.reshape(centroids, [self.B*self.nb_clusters, self.E])

				# Add + [0,1,2]
				idx_flattened = tf.reshape(labels + self.shifting, [self.B*self.L]) # [b*nb_tries*L]
				X_flattened = tf.reshape(self.X, [self.B*self.L, self.E]) # [b*nb_tries*L, E]

				dist = tf.reduce_sum(tf.square(X_flattened - tf.gather(centroids_flattened, idx_flattened)), -1)

				total = tf.unsorted_segment_sum(dist, idx_flattened, self.B*self.nb_clusters)
				count = tf.unsorted_segment_sum(tf.ones_like(dist), idx_flattened, self.B*self.nb_clusters)

				# [self.B*self.nb_clusters]
				inertia = total / count
				inertia = tf.reshape(inertia, [self.B, self.nb_clusters])
				inertia = tf.reduce_sum(inertia, -1)

			return inertia

	def body(self ,i, centroids, labels):
		with tf.name_scope('iteration'):

			X = self.X * self.notsilent

			# Add + [0,1,2]
			if self.beta is not None:
				X_ = tf.expand_dims(X, 2) # B L 1 E
				labels_ = tf.expand_dims(labels, 3) # B L C 1
				new_centroids = tf.reduce_sum(X_ * labels_, 1) / tf.expand_dims(tf.reduce_sum(labels, 1),-1)
				new_centroids = tf.reshape(new_centroids, [self.B, self.nb_clusters, self.E])

			else:
				labels_flattened = tf.reshape(labels + self.shifting, [-1])
				X_flattened = tf.reshape(X, [-1, self.E])

				total = tf.unsorted_segment_sum(X_flattened, labels_flattened, self.B*self.nb_clusters)
				count = tf.unsorted_segment_sum(tf.ones_like(X_flattened), labels_flattened, self.B*self.nb_clusters)

				new_centroids = total/count
				new_centroids = tf.reshape(new_centroids, [self.B, self.nb_clusters, self.E])

			return [i+1, new_centroids, self.get_labels(new_centroids)]

	def get_labels(self, centroids):

		X = self.X

		X_ = tf.expand_dims(X, 2) # B L 1 E
		centroids_ = tf.expand_dims(centroids, 1) # B 1 C E

		if self.beta is not None:
			distances = tf.reduce_sum(tf.square(X_ - centroids_) * tf.expand_dims(self.notsilent, 2), axis=3)
			exp = tf.exp(-1.0 *  self.beta * distances)
			exp =  exp / tf.reduce_sum(exp, -1, keep_dims=True)
			return exp
		else:
			distances = tf.sqrt(tf.reduce_sum(tf.square(X_ - centroids_) * tf.expand_dims(self.notsilent, 2), axis=3))
			return tf.argmin(distances, axis=2, output_type=tf.int32)

	def fit(self, X_train):
		with tf.Session(graph=self.graph) as sess:
			sess.run(tf.global_variables_initializer())
			return sess.run(self.network, {self.X_in: X_train})

# from sklearn.cluster import KMeans as km

if __name__ == "__main__":
	nb_samples = 1000
	E = 40
	nb_clusters = 4
	error = 0
	TOTAL = 1
	nb_err = 0.0
	kmean = KMeans(nb_clusters, nb_tries=10, nb_iterations=10, beta=0.5)

	for i in range(TOTAL):
		X, y = make_blobs(n_samples=nb_samples, centers=nb_clusters, n_features=E, cluster_std=1.0)
		X_ = X[np.newaxis,:]
		X_ = np.concatenate([X_, X_, X_], axis=0)
		y = y[np.newaxis,:]

		# with tf.Session(graph=kmean.graph) as sess:
		# 	sess.run(tf.global_variables_initializer())
		# 	print sess.run(kmean.network, {kmean.X_in: X_})

		centroids, labels = kmean.fit(X_)
		print labels
		print np.argmax(labels, 2)
		print y
		# print inertia

		# print labels
		# u = tf.one_hot(labels, 2, 1.0, 0.0, name='masks')
		# masked_val = X_ * u
		# t = tf.transpose(masked_val, [0,2,1])
		# mixed = tf.reshape(t , [2*nb_clusters, -1])
		# # print masked_val
		# # with tf.Session().as_default():
		# # 	print mixed.eval()

		# centroids = np.reshape(centroids, (2, nb_clusters, E))
		# kmeans = km(n_clusters=nb_clusters, random_state=0, ).fit(X)
<commit_msg>Add 'assign end' feature to KMEANS<commit_after># -*- coding: utf-8 -*-
# My Model
from utils.ops import scope
import tensorflow as tf
import numpy as np
from  sklearn.datasets import make_blobs

#############################################
#       Deep Adaptive Separator Model       #
#############################################

class KMeans:

	def __init__(self, nb_clusters, centroids_init=None, nb_tries=10,
		nb_iterations=10, input_tensor=None, latent_space_tensor=None, beta=None, threshold=2.5, assign_at_end=False):

		self.nb_clusters = nb_clusters
		self.nb_iterations = nb_iterations
		self.nb_tries = nb_tries
		self.latent_space_tensor = latent_space_tensor
		self.beta = beta
		self.assign_at_end = assign_at_end

		if input_tensor is None:
			self.graph = tf.Graph()
		else:
			self.graph = tf.get_default_graph()

		with self.graph.as_default():
			with tf.name_scope('kmeans'):
				if input_tensor is None:
					# Spectrogram, embeddings
					# shape = [batch, L , E ]
					self.X_in = tf.placeholder("float", [None, None, None], name='Kmeans_input')
				else:
					self.X_in = input_tensor

				# mean, _ = tf.nn.moments(self.X_in, axes=-1, keep_dims=True)
				x_norm = tf.nn.l2_normalize(self.X_in, axis=-1)
				self.b = tf.shape(x_norm)[0]
				self.X_original = tf.identity(x_norm)
				self.X = tf.expand_dims(x_norm, 1)
				self.X = tf.tile(self.X, [1, self.nb_tries, 1, 1])

				self.L = tf.shape(self.X)[-2]
				self.E = tf.shape(self.X)[-1]
				self.X = tf.reshape(self.X, [-1, self.L, self.E])

				self.B = tf.shape(self.X)[0]

				self.ones = tf.ones_like(self.X, tf.float32)

				self.shifting = tf.tile(tf.expand_dims(tf.range(self.B)*self.nb_clusters, 1), [1, self.L])

				if centroids_init is None:
					def random_without_replace(b, l):
						a =  np.array([np.random.choice(range(l), size=self.nb_clusters, replace=False) for _ in range(b)])
						return a.astype(np.int32)

					y = tf.py_func(random_without_replace, [self.B, self.L], tf.int32)
					random = tf.reshape(y, [self.B, self.nb_clusters, 1])

					# Take randomly 'nb_clusters' vectors from X
					batch_range = tf.tile(tf.reshape(tf.range(self.B, dtype=tf.int32), shape=[self.B, 1, 1]), [1, self.nb_clusters, 1])
					indices = tf.concat([batch_range, random], axis = 2)
					self.centroid_init = tf.gather_nd(self.X, indices)
				else:
					self.centroids = tf.identity(centroids_init)
					self.centroids = tf.tile(self.centroids, [self.nb_tries, 1 , 1])

				if not self.latent_space_tensor is None:
					latent_space_tensor = tf.reshape(latent_space_tensor, [self.b, self.L])
					log_lst = tf.log(tf.divide(tf.reduce_max(latent_space_tensor, [-1], keep_dims=True),latent_space_tensor))
					self.notsilent_notry = tf.reshape(tf.cast(log_lst < threshold, tf.float32), [self.b, self.L, 1])
					self.notsilent = tf.tile(self.notsilent_notry, [self.nb_tries, 1 , 1])
				else:
					self.notsilent = tf.ones([self.B, self.L, 1])

				self.network

	@scope
	def network(self):
		i = tf.constant(0)
		init = [i, self.centroid_init, self.get_labels(self.centroid_init)]
		if self.beta is not None:
			invariant = [i.get_shape(), tf.TensorShape([None, self.nb_clusters, None]), tf.TensorShape([None, None, self.nb_clusters])]
		else:
			invariant = [i.get_shape(), tf.TensorShape([None, self.nb_clusters, None]), tf.TensorShape([None, None])]

		cond = lambda i, c, l: tf.less(i, self.nb_iterations)
		_ , centroids, labels = tf.while_loop(cond, self.body, init, shape_invariants=invariant)
		inertia = self.get_inertia(centroids)

		inertia = tf.reshape(inertia, [self.b, self.nb_tries])
		bests = tf.argmin(inertia, 1, output_type=tf.int32)
		index = bests + tf.range(self.b)*self.nb_tries

		centroids = tf.reshape(centroids, [self.b*self.nb_tries, self.nb_clusters, self.E])
		centroids = tf.gather(centroids, index)

		if self.assign_at_end:
			labels = self.get_labels(centroids, end=True)
		else:
			labels = tf.gather(labels, index)

		return centroids, labels

	def get_inertia(self, centroids):
		with tf.name_scope('inertia'):
			labels = self.get_labels(centroids)
			if self.beta is not None:
				X_ = tf.expand_dims(self.X, 2) # B L 1 E
				centroids_ = tf.expand_dims(centroids, 1) # B 1 C E
				dist = tf.reduce_sum(tf.square(X_ - centroids_), -1) * labels # B C distance to clusters according to soft assignments
				dist = tf.reduce_sum(dist, 1) # Sum over points
				density = tf.reduce_sum(labels, 1) # Compute the soft density of each cluster
				inertia = tf.divide(dist, density)
				inertia = tf.reduce_sum(inertia, -1)
			else:
				# centroids [b*nb_tries, C, E]
				centroids_flattened = tf.reshape(centroids, [self.B*self.nb_clusters, self.E])

				# Add + [0,1,2]
				idx_flattened = tf.reshape(labels + self.shifting, [self.B*self.L]) # [b*nb_tries*L]
				X_flattened = tf.reshape(self.X, [self.B*self.L, self.E]) # [b*nb_tries*L, E]

				dist = tf.reduce_sum(tf.square(X_flattened - tf.gather(centroids_flattened, idx_flattened)), -1)

				total = tf.unsorted_segment_sum(dist, idx_flattened, self.B*self.nb_clusters)
				count = tf.unsorted_segment_sum(tf.ones_like(dist), idx_flattened, self.B*self.nb_clusters)

				# [self.B*self.nb_clusters]
				inertia = total / count
				inertia = tf.reshape(inertia, [self.B, self.nb_clusters])
				inertia = tf.reduce_sum(inertia, -1)

			return inertia

	def body(self ,i, centroids, labels):
		with tf.name_scope('iteration'):

			X = self.X * self.notsilent

			# Add + [0,1,2]
			if self.beta is not None:
				X_ = tf.expand_dims(X, 2) # B L 1 E
				labels_ = tf.expand_dims(labels, 3) # B L C 1
				new_centroids = tf.reduce_sum(X_ * labels_, 1) / tf.expand_dims(tf.reduce_sum(labels, 1),-1)
				new_centroids = tf.reshape(new_centroids, [self.B, self.nb_clusters, self.E])

			else:
				labels_flattened = tf.reshape(labels + self.shifting, [-1])
				X_flattened = tf.reshape(X, [-1, self.E])

				total = tf.unsorted_segment_sum(X_flattened, labels_flattened, self.B*self.nb_clusters)
				count = tf.unsorted_segment_sum(tf.ones_like(X_flattened), labels_flattened, self.B*self.nb_clusters)

				new_centroids = total/count
				new_centroids = tf.reshape(new_centroids, [self.B, self.nb_clusters, self.E])

			return [i+1, new_centroids, self.get_labels(new_centroids)]

	def get_labels(self, centroids, end=False):

		if end:
			notsilent = tf.ones([self.b, self.L, 1])
			X = self.X_original
		else:
			notsilent = self.notsilent
			X = self.X

		X_ = tf.expand_dims(X, 2) # B L 1 E
		centroids_ = tf.expand_dims(centroids, 1) # B 1 C E

		if self.beta is not None:
			distances = tf.reduce_sum(tf.square(X_ - centroids_) * tf.expand_dims(notsilent, 2), axis=3)
			exp = tf.exp(-1.0 *  self.beta * distances)
			exp =  exp / tf.reduce_sum(exp, -1, keep_dims=True)
			return exp
		else:
			distances = tf.sqrt(tf.reduce_sum(tf.square(X_ - centroids_) * tf.expand_dims(notsilent, 2), axis=3))
			return tf.argmin(distances, axis=2, output_type=tf.int32)

	def fit(self, X_train):
		with tf.Session(graph=self.graph) as sess:
			sess.run(tf.global_variables_initializer())
			return sess.run(self.network, {self.X_in: X_train})

# from sklearn.cluster import KMeans as km

if __name__ == "__main__":
	nb_samples = 1000
	E = 40
	nb_clusters = 4
	error = 0
	TOTAL = 1
	nb_err = 0.0
	kmean = KMeans(nb_clusters, nb_tries=10, nb_iterations=10, beta=0.5)

	for i in range(TOTAL):
		X, y = make_blobs(n_samples=nb_samples, centers=nb_clusters, n_features=E, cluster_std=1.0)
		X_ = X[np.newaxis,:]
		X_ = np.concatenate([X_, X_, X_], axis=0)
		y = y[np.newaxis,:]

		# with tf.Session(graph=kmean.graph) as sess:
		# 	sess.run(tf.global_variables_initializer())
		# 	print sess.run(kmean.network, {kmean.X_in: X_})

		centroids, labels = kmean.fit(X_)
		print labels
		print np.argmax(labels, 2)
		print y
		# print inertia

		# print labels
		# u = tf.one_hot(labels, 2, 1.0, 0.0, name='masks')
		# masked_val = X_ * u
		# t = tf.transpose(masked_val, [0,2,1])
		# mixed = tf.reshape(t , [2*nb_clusters, -1])
		# # print masked_val
		# # with tf.Session().as_default():
		# # 	print mixed.eval()

		# centroids = np.reshape(centroids, (2, nb_clusters, E))
		# kmeans = km(n_clusters=nb_clusters, random_state=0, ).fit(X)
<|endoftext|>


<commit_before>import util.main as main
from wodpy import wod

class TestClass:
   def setUp(self):
        filenames = main.readInput('datafiles.json')
        profiles = main.extractProfiles(filenames)

        # identify and import tests
        testNames = main.importQC('qctests')
        testNames.sort()
        for testName in testNames:
          exec('from qctests import ' + testName)

        # Set up any keyword arguments needed by tests.
        kwargs = {'profiles' : profiles}

        testResults  = []
        testVerbose  = []
        trueResults  = []
        trueVerbose  = []
        firstProfile = True
        delete       = []
        currentFile  = ''
        self.profiles = []
        for iprofile, pinfo in enumerate(profiles):
          # Load the profile data.
          if pinfo.file_name != currentFile:
            if currentFile != '': f.close()
            currentFile = pinfo.file_name
            f = open(currentFile)
          if f.tell() != pinfo.file_position: f.seek(pinfo.file_position)
          self.profiles.append(wod.WodProfile(f))

   def tearDown(self):
       return

   def test_check_time_type(self):
        for p in self.profiles:
            time = p.time()
            assert type(time) is float or time is None, 'profile time should be either a float or None.'  <commit_msg>stopping wod test from breaking when datafiles.json changes<commit_after>import util.main as main
from wodpy import wod

class TestClass:
   def setUp(self):
        filenames = ["data/quota_subset.dat"]
        profiles = main.extractProfiles(filenames)

        # identify and import tests
        testNames = main.importQC('qctests')
        testNames.sort()
        for testName in testNames:
          exec('from qctests import ' + testName)

        # Set up any keyword arguments needed by tests.
        kwargs = {'profiles' : profiles}

        testResults  = []
        testVerbose  = []
        trueResults  = []
        trueVerbose  = []
        firstProfile = True
        delete       = []
        currentFile  = ''
        self.profiles = []
        for iprofile, pinfo in enumerate(profiles):
          # Load the profile data.
          if pinfo.file_name != currentFile:
            if currentFile != '': f.close()
            currentFile = pinfo.file_name
            f = open(currentFile)
          if f.tell() != pinfo.file_position: f.seek(pinfo.file_position)
          self.profiles.append(wod.WodProfile(f))

   def tearDown(self):
       return

   def test_check_time_type(self):
        for p in self.profiles:
            time = p.time()
            assert type(time) is float or time is None, 'profile time should be either a float or None.'  <|endoftext|>


<commit_before>import logging
from flask import jsonify, request
import flask_login

from server import app, mc
from server.auth import user_mediacloud_client
from server.util.request import form_fields_required, api_error_handler, json_error_response, arguments_required
from server.util.stringutil import ids_from_comma_separated_str
from server.util.tags import US_COLLECTIONS
from server.views.topics import concatenate_query_for_solr, concatenate_solr_dates
from server.views.topics.foci.retweetpartisanship import add_retweet_partisanship_to_topic
from server.views.topics.topic import topic_summary

logger = logging.getLogger(__name__)
VERSION_1 = 1

@app.route('/api/topics/create/preview/split-story/count', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_split_story_count():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,
    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    results = user_mc.storyCount(solr_query=solr_query, solr_filter=fq, split=True)
    total_stories = 0
    for c in results['counts']:
        total_stories += c['count']
    results['total_story_count'] = total_stories
    return jsonify({'results': results})

@app.route('/api/topics/create/preview/story/count', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_story_count():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,
    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    story_count_result = user_mc.storyCount(solr_query=solr_query, solr_filter=fq)
    # maybe check admin role before we run this?
    return jsonify(story_count_result)  # give them back new data, so they can update the client

@app.route('/api/topics/create/preview/stories/sample', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_story_sample():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,

    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    num_stories = request.form['rows']
    story_count_result = user_mc.storyList(solr_query=solr_query, solr_filter=fq, sort=user_mc.SORT_RANDOM, rows=num_stories)
    return jsonify(story_count_result)

@app.route('/api/topics/create/preview/words/count', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_word_count():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,
    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    word_count_result = user_mc.wordCount(solr_query=solr_query, solr_filter=fq)

    return jsonify(word_count_result)  # give them back new data, so they can update the client

@app.route('/api/topics/create', methods=['PUT'])
@flask_login.login_required
@form_fields_required('name', 'description', 'solr_seed_query', 'start_date', 'end_date')
@api_error_handler
def topic_create():
    user_mc = user_mediacloud_client()
    name = request.form['name']
    description = request.form['description']
    solr_seed_query = request.form['solr_seed_query']
    start_date = request.form['start_date']
    end_date = request.form['end_date']

    optional_args = {
        'is_public': request.form['is_public'] if 'is_public' in request.form else None,
        'is_logogram': request.form['is_logogram'] if 'is_logogram' in request.form else None,
        'ch_monitor_id': request.form['ch_monitor_id'] if len(request.form['ch_monitor_id']) > 0 and request.form['ch_monitor_id'] != 'null' else None,
        'max_iterations': request.form['max_iterations'] if 'max_iterations' in request.form else None,
        'max_stories': request.form['max_stories'] if 'max_stories' in request.form and request.form['max_stories'] != 'null' else flask_login.current_user.profile['max_topic_stories'],
    }

    # parse out any sources and collections to add
    media_ids_to_add = ids_from_comma_separated_str(request.form['sources[]'])
    tag_ids_to_add = ids_from_comma_separated_str(request.form['collections[]'])

    try:
        topic_result = user_mc.topicCreate(name=name, description=description, solr_seed_query=solr_seed_query,
                                           start_date=start_date, end_date=end_date, media_ids=media_ids_to_add,
                                           media_tags_ids=tag_ids_to_add, **optional_args)['topics'][0]

        topics_id = topic_result['topics_id']
        logger.info("Created new topic \"{}\" as {}".format(name, topics_id))
        # if this includes any of the US-centric collections, add the retweet partisanship subtopic by default
        if set(tag_ids_to_add).intersection(US_COLLECTIONS):
            add_retweet_partisanship_to_topic(topic_result['topics_id'],
                                              'Retweet Partisanship',
                                              'Subtopics driven by our analysis of Twitter followers of Trump and Clinton during the 2016 election season.  Each media soure is scored based on the ratio of retweets of their stories in those two groups.')
        # client will either make a empty snapshot, or a spidering one
        return topic_summary(topics_id)
    except Exception as e:
        logging.error("Topic creation failed {}".format(name))
        logging.exception(e)
        return json_error_response(e.message, e.status_code)

@app.route('/api/topics/name-exists', methods=['GET'])
@flask_login.login_required
@arguments_required('searchStr')
@api_error_handler
def topic_name_exists():
    # Check if topic with name exists already
    # Have to do this in a unique method, instead of in topic_search because we need to use an admin connection
    # to media cloud to list all topics, but we don't want to return topics a user can't see to them.
    # :return: boolean indicating if topic with this name exists for not (case insensive check)
    search_str = request.args['searchStr']
    topics_id = int(request.args['topicId']) if 'topicId' in request.args else None
    matching_topics = mc.topicList(name=search_str, limit=15)
    if topics_id:
        matching_topic_names = [t['name'].lower().strip() for t in matching_topics['topics']
                                if t['topics_id'] != topics_id]
    else:
        matching_topic_names = [t['name'].lower().strip() for t in matching_topics['topics']]
    name_in_use = search_str.lower() in matching_topic_names
    return jsonify({'nameInUse': name_in_use})
<commit_msg>better error handlingn<commit_after>import logging
from flask import jsonify, request
import flask_login
import mediacloud.error

from server import app, mc
from server.auth import user_mediacloud_client
from server.util.request import form_fields_required, api_error_handler, json_error_response, arguments_required
from server.util.stringutil import ids_from_comma_separated_str
from server.util.tags import US_COLLECTIONS
from server.views.topics import concatenate_query_for_solr, concatenate_solr_dates
from server.views.topics.foci.retweetpartisanship import add_retweet_partisanship_to_topic
from server.views.topics.topic import topic_summary

logger = logging.getLogger(__name__)
VERSION_1 = 1

@app.route('/api/topics/create/preview/split-story/count', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_split_story_count():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,
    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    results = user_mc.storyCount(solr_query=solr_query, solr_filter=fq, split=True)
    total_stories = 0
    for c in results['counts']:
        total_stories += c['count']
    results['total_story_count'] = total_stories
    return jsonify({'results': results})

@app.route('/api/topics/create/preview/story/count', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_story_count():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,
    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    story_count_result = user_mc.storyCount(solr_query=solr_query, solr_filter=fq)
    # maybe check admin role before we run this?
    return jsonify(story_count_result)  # give them back new data, so they can update the client

@app.route('/api/topics/create/preview/stories/sample', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_story_sample():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,

    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    num_stories = request.form['rows']
    story_count_result = user_mc.storyList(solr_query=solr_query, solr_filter=fq, sort=user_mc.SORT_RANDOM, rows=num_stories)
    return jsonify(story_count_result)

@app.route('/api/topics/create/preview/words/count', methods=['POST'])
@flask_login.login_required
@form_fields_required('q')
@api_error_handler
def api_topics_preview_word_count():
    user_mc = user_mediacloud_client()

    solr_query = concatenate_query_for_solr(solr_seed_query=request.form['q'],
                                            media_ids=ids_from_comma_separated_str(request.form['sources[]']) if 'sources[]' in request.form else None,
                                            tags_ids=ids_from_comma_separated_str(request.form['collections[]'])) if 'collections[]' in request.form else None,
    fq = concatenate_solr_dates(start_date=request.form['start_date'],
                                end_date=request.form['end_date'])
    word_count_result = user_mc.wordCount(solr_query=solr_query, solr_filter=fq)

    return jsonify(word_count_result)  # give them back new data, so they can update the client

@app.route('/api/topics/create', methods=['PUT'])
@flask_login.login_required
@form_fields_required('name', 'description', 'solr_seed_query', 'start_date', 'end_date')
@api_error_handler
def topic_create():
    user_mc = user_mediacloud_client()
    name = request.form['name']
    description = request.form['description']
    solr_seed_query = request.form['solr_seed_query']
    start_date = request.form['start_date']
    end_date = request.form['end_date']

    optional_args = {
        'is_public': request.form['is_public'] if 'is_public' in request.form else None,
        'is_logogram': request.form['is_logogram'] if 'is_logogram' in request.form else None,
        'ch_monitor_id': request.form['ch_monitor_id'] if len(request.form['ch_monitor_id']) > 0 and request.form['ch_monitor_id'] != 'null' else None,
        'max_iterations': request.form['max_iterations'] if 'max_iterations' in request.form else None,
        'max_stories': request.form['max_stories'] if 'max_stories' in request.form and request.form['max_stories'] != 'null' else flask_login.current_user.profile['max_topic_stories'],
    }

    # parse out any sources and collections to add
    media_ids_to_add = ids_from_comma_separated_str(request.form['sources[]'])
    tag_ids_to_add = ids_from_comma_separated_str(request.form['collections[]'])

    try:
        topic_result = user_mc.topicCreate(name=name, description=description, solr_seed_query=solr_seed_query,
                                           start_date=start_date, end_date=end_date, media_ids=media_ids_to_add,
                                           media_tags_ids=tag_ids_to_add, **optional_args)['topics'][0]

        topics_id = topic_result['topics_id']
        logger.info("Created new topic \"{}\" as {}".format(name, topics_id))
        # if this includes any of the US-centric collections, add the retweet partisanship subtopic by default
        if set(tag_ids_to_add).intersection(US_COLLECTIONS):
            add_retweet_partisanship_to_topic(topic_result['topics_id'],
                                              'Retweet Partisanship',
                                              'Subtopics driven by our analysis of Twitter followers of Trump and Clinton during the 2016 election season.  Each media soure is scored based on the ratio of retweets of their stories in those two groups.')
        # client will either make a empty snapshot, or a spidering one
        return topic_summary(topics_id)
    except Exception as e:
        logging.error("Topic creation failed {}".format(name))
        logging.exception(e)
        return json_error_response(str(e), 500)
    except mediacloud.error.MCException as e:
        logging.error("Topic creation failed {}".format(name))
        logging.exception(e)
        return json_error_response(e.message, e.status_code)

@app.route('/api/topics/name-exists', methods=['GET'])
@flask_login.login_required
@arguments_required('searchStr')
@api_error_handler
def topic_name_exists():
    # Check if topic with name exists already
    # Have to do this in a unique method, instead of in topic_search because we need to use an admin connection
    # to media cloud to list all topics, but we don't want to return topics a user can't see to them.
    # :return: boolean indicating if topic with this name exists for not (case insensive check)
    search_str = request.args['searchStr']
    topics_id = int(request.args['topicId']) if 'topicId' in request.args else None
    matching_topics = mc.topicList(name=search_str, limit=15)
    if topics_id:
        matching_topic_names = [t['name'].lower().strip() for t in matching_topics['topics']
                                if t['topics_id'] != topics_id]
    else:
        matching_topic_names = [t['name'].lower().strip() for t in matching_topics['topics']]
    name_in_use = search_str.lower() in matching_topic_names
    return jsonify({'nameInUse': name_in_use})
<|endoftext|>


<commit_before>from Q50_config import *
import sys, os
from GPSReader import *
from GPSReprojection import *
from GPSTransforms import *
from VideoReader import *
from LidarTransforms import *
from ColorMap import *
from transformations import euler_matrix
import numpy as np
import cv2
from ArgParser import *

WINDOW = 50*5

def cloudToPixels(cam, pts_wrt_cam):

    width = 4
    (pix, J)  = cv2.projectPoints(pts_wrt_cam.transpose(), np.array([0.0,0.0,0.0]), np.array([0.0,0.0,0.0]), cam['KK'], cam['distort'])

    pix = pix.transpose()
    pix = np.around(pix[:, 0, :])
    pix = pix.astype(np.int32)
    mask = np.logical_and(True, pix[0,:] > 0 + width/2)
    mask = np.logical_and(mask, pix[1,:] > 0 + width/2)
    mask = np.logical_and(mask, pix[0,:] < 1279 - width/2)
    mask = np.logical_and(mask, pix[1,:] < 959 - width/2)
    mask = np.logical_and(mask, pts_wrt_cam[2,:] > 0)
    dist_sqr = np.sum( pts_wrt_cam[0:3, :] ** 2, axis = 0)
    mask = np.logical_and(mask, dist_sqr > 3)

    return (pix, mask)

def localMapToPixels(map_data, imu_transforms_t, T_from_i_to_l, cam):
    # load nearby map frames
    pts_wrt_imu_0 = array(map_data[:,0:3]).transpose()
    pts_wrt_imu_0 = np.vstack((pts_wrt_imu_0,
        np.ones((1,pts_wrt_imu_0.shape[1]))))
    # transform points from imu_0 to imu_t
    pts_wrt_imu_t = np.dot( np.linalg.inv(imu_transforms_t), pts_wrt_imu_0)
    # transform points from imu_t to lidar_t
    pts_wrt_lidar_t = np.dot(T_from_i_to_l, pts_wrt_imu_t);
    # transform points from lidar_t to camera_t
    pts_wrt_camera_t = pts_wrt_lidar_t.transpose()[:, 0:3] + cam['displacement_from_l_to_c_in_lidar_frame']
    pts_wrt_camera_t = dot(R_to_c_from_l(cam),
            pts_wrt_camera_t.transpose())
    # reproject camera_t points in camera frame
    (pix, mask) = cloudToPixels(cam, pts_wrt_camera_t)

    return (pix, mask)

if __name__ == '__main__':
    args = parse_args(sys.argv[1], sys.argv[2])
    cam_num = int(sys.argv[2][-5])

    params = LoadParameters('q50_4_3_14_params')
    cam = params['cam'][cam_num-1]
    video_reader = VideoReader(args['video'])
    gps_reader = GPSReader(args['gps'])
    GPSData = gps_reader.getNumericData()
    imu_transforms = IMUTransforms(GPSData)

    T_from_i_to_l = np.linalg.inv(params['lidar']['T_from_l_to_i'])

    all_data = np.load(sys.argv[3])
    map_data = all_data['data']
    #map_data = map_data[map_data[:,3] > 60, :]
    # map points are defined w.r.t the IMU position at time 0
    # each entry in map_data is (x,y,z,intensity,framenum).

    print "Hit 'q' to quit"
    while True:
        for count in range(20):
            (success, I) = video_reader.getNextFrame()

        t = video_reader.framenum - 1
        print t
        mask_window = (map_data[:,4] < t + WINDOW) & (map_data[:,4] > t );
        map_data_copy = array(map_data[mask_window, :]);

        # reproject
        (pix, mask) = localMapToPixels(map_data_copy, imu_transforms[t,:,:], T_from_i_to_l, cam);

        # draw
        intensity = map_data_copy[mask, 3]
        heat_colors = heatColorMapFast(intensity, 0, 100)
        for p in range(4):
            I[pix[1,mask]+p, pix[0,mask], :] = heat_colors[0,:,:]
            I[pix[1,mask], pix[0,mask]+p, :] = heat_colors[0,:,:]
            I[pix[1,mask]+p, pix[0,mask], :] = heat_colors[0,:,:]
            I[pix[1,mask], pix[0,mask]+p, :] = heat_colors[0,:,:]

        cv2.imshow('vid', cv2.pyrDown(I))
        keycode = cv2.waitKey(1)
        if keycode == 113:
            break

<commit_msg>break when no more video frames<commit_after>from Q50_config import *
import sys, os
from GPSReader import *
from GPSReprojection import *
from GPSTransforms import *
from VideoReader import *
from LidarTransforms import *
from ColorMap import *
from transformations import euler_matrix
import numpy as np
import cv2
from ArgParser import *

WINDOW = 50*5

def cloudToPixels(cam, pts_wrt_cam):

    width = 4
    (pix, J)  = cv2.projectPoints(pts_wrt_cam.transpose(), np.array([0.0,0.0,0.0]), np.array([0.0,0.0,0.0]), cam['KK'], cam['distort'])

    pix = pix.transpose()
    pix = np.around(pix[:, 0, :])
    pix = pix.astype(np.int32)
    mask = np.logical_and(True, pix[0,:] > 0 + width/2)
    mask = np.logical_and(mask, pix[1,:] > 0 + width/2)
    mask = np.logical_and(mask, pix[0,:] < 1279 - width/2)
    mask = np.logical_and(mask, pix[1,:] < 959 - width/2)
    mask = np.logical_and(mask, pts_wrt_cam[2,:] > 0)
    dist_sqr = np.sum( pts_wrt_cam[0:3, :] ** 2, axis = 0)
    mask = np.logical_and(mask, dist_sqr > 3)

    return (pix, mask)

def localMapToPixels(map_data, imu_transforms_t, T_from_i_to_l, cam):
    # load nearby map frames
    pts_wrt_imu_0 = array(map_data[:,0:3]).transpose()
    pts_wrt_imu_0 = np.vstack((pts_wrt_imu_0,
        np.ones((1,pts_wrt_imu_0.shape[1]))))
    # transform points from imu_0 to imu_t
    pts_wrt_imu_t = np.dot( np.linalg.inv(imu_transforms_t), pts_wrt_imu_0)
    # transform points from imu_t to lidar_t
    pts_wrt_lidar_t = np.dot(T_from_i_to_l, pts_wrt_imu_t);
    # transform points from lidar_t to camera_t
    pts_wrt_camera_t = pts_wrt_lidar_t.transpose()[:, 0:3] + cam['displacement_from_l_to_c_in_lidar_frame']
    pts_wrt_camera_t = dot(R_to_c_from_l(cam),
            pts_wrt_camera_t.transpose())
    # reproject camera_t points in camera frame
    (pix, mask) = cloudToPixels(cam, pts_wrt_camera_t)

    return (pix, mask)

if __name__ == '__main__':
    args = parse_args(sys.argv[1], sys.argv[2])
    cam_num = int(sys.argv[2][-5])

    params = LoadParameters('q50_4_3_14_params')
    cam = params['cam'][cam_num-1]
    video_reader = VideoReader(args['video'])
    gps_reader = GPSReader(args['gps'])
    GPSData = gps_reader.getNumericData()
    imu_transforms = IMUTransforms(GPSData)

    T_from_i_to_l = np.linalg.inv(params['lidar']['T_from_l_to_i'])

    all_data = np.load(sys.argv[3])
    map_data = all_data['data']
    #map_data = map_data[map_data[:,3] > 60, :]
    # map points are defined w.r.t the IMU position at time 0
    # each entry in map_data is (x,y,z,intensity,framenum).

    print "Hit 'q' to quit"
    while True:
        for count in range(20):
            (success, I) = video_reader.getNextFrame()

        if not success:
            break

        t = video_reader.framenum - 1
        print t
        mask_window = (map_data[:,4] < t + WINDOW) & (map_data[:,4] > t );
        map_data_copy = array(map_data[mask_window, :]);

        # reproject
        (pix, mask) = localMapToPixels(map_data_copy, imu_transforms[t,:,:], T_from_i_to_l, cam);

        # draw
        intensity = map_data_copy[mask, 3]
        heat_colors = heatColorMapFast(intensity, 0, 100)
        for p in range(4):
            I[pix[1,mask]+p, pix[0,mask], :] = heat_colors[0,:,:]
            I[pix[1,mask], pix[0,mask]+p, :] = heat_colors[0,:,:]
            I[pix[1,mask]+p, pix[0,mask], :] = heat_colors[0,:,:]
            I[pix[1,mask], pix[0,mask]+p, :] = heat_colors[0,:,:]

        cv2.imshow('vid', cv2.pyrDown(I))
        keycode = cv2.waitKey(1)
        if keycode == 113:
            break

<|endoftext|>


<commit_before>import bitstring, serial, struct
from collections import namedtuple, OrderedDict
from datetime import datetime, timedelta

SerialPayload = namedtuple('SerialPayload', [ 'device_id', 'payload' ])
GenericCANMsg = namedtuple('GenericCANMsg', [ 'device_type', 'manufacturer',
                                              'api_class', 'api_key',
                                              'device_number', 'payload' ])
MotorCtrl     = namedtuple('MotorCtrl',     [ 'device_type', 'manufacturer',
                                              'device_number' ])

class Packetizer:
    def __init__(self, sof, esc, sof_esc, esc_esc):
        self.sof = bytearray([ sof ])
        self.esc = bytearray([ esc ])
        self.esc_sof = bytearray([ sof_esc ])
        self.esc_esc = bytearray([ esc_esc ])
        self._reset()

    def recv_byte(self, byte):
        byte_dec = self._decode(self.last, byte)

        # Every start of frame (SOF) byte starts a new frame because it is
        # otherwise escaped.
        if byte == self.sof:
            self._reset()
        # Next byte is the total number of bytes in the packet. Note that this
        # could be encoded if it equals 255, although this should never occur
        # in practice when framing CANbus messages.
        elif self.offset == 1:
            self.count   = ord(byte_dec)
            self.offset += 1
        elif byte_dec:
            self.packet.extend(byte_dec)
            self.offset += 1

        self.last = byte

        # Done receiving the entire packet, so add it to the queue. Padding
        # compensates for fields in the packet (e.g. device id) that don't
        # contribute to the byte count.
        if self.valid and self.offset == self.count + 2:
            new_packet = self.packet
            self._reset()
            return new_packet
        else:
            return None

    def frame_bytes(self, payload):
        count = len(payload)
        contents = struct.pack('B{}s'.format(count), count, payload)
        return self.sof + self._encode(contents)

    def _reset(self):
        self.offset = 1
        self.count  = 0
        self.last   = None
        self.packet = bytearray()
        self.valid  = True

    def _error(self):
        self.valid = False

    def _decode(self, last, curr):
        if last == self.esc:
            if curr == self.esc_esc:
                return self.esc
            elif curr == self.esc_sof:
                return self.sof
        elif curr == self.esc:
            return None
        else:
            return curr

    def _encode(self, raw):
        encoded = bytearray()
        for curr in raw:
            if curr == self.sof:
                encoded.extend(self.esc)
                encoded.extend(self.esc_sof)
            elif curr == self.esc:
                encoded.extend(self.esc)
                encoded.extend(self.esc_esc)
            else:
                encoded.extend(curr)
        return encoded

class JaguarUART:
    header_fields = OrderedDict([
        ('resv',          'uint:3'),
        ('device_type',   'uint:5'),
        ('manufacturer',  'uint:8'),
        ('api_class',     'uint:6'),
        ('api_key',       'uint:4'),
        ('device_number', 'uint:6')
    ])
    header_fmt = [ field + '=' + name for name, field in header_fields.items() ]

    def __init__(self, serial, packetizer):
        self.serial = serial
        self.packetizer = packetizer
        self.alive = True

    def close(self):
        self.alive = False

    @classmethod
    def parse_message(cls, packet):
        header_raw = bitstring.BitArray(bytes=packet[0:4])
        header_raw.byteswap()

        header_list = header_raw.unpack(cls.header_fmt)
        header = dict(zip(cls.header_fields.keys(), header_list))
        del header['resv']
        return GenericCANMsg(payload=packet[4:], **header)

    @classmethod
    def generate_message(cls, msg):
        header = bitstring.pack(cls.header_fmt, resv=0, **msg._asdict())
        header.byteswap()
        return header.bytes + msg.payload

    def send_message(self, msg):
        packed = self.generate_message(msg)
        framed = self.packetizer.frame_bytes(packed)
        self.serial.write(framed)

    def recv_message(self, timeout=None):
        self.timeout = timeout
        while True:
            # This may timeout and return zero bytes, so timeout here as well.
            byte = self.serial.read(1)
            if byte:
                return None

            # Feed the packetizing state machine the next character.
            packet = self.packetizer.recv_byte(byte)
            if packet:
                return self.parse_message(packet)

    def pending_data(self):
        return self.serial.inWaiting() > 0

class JaguarCAN:
    class DeviceType:
        BroadcastMessages   = 0
        RobotController     = 1
        MotorController     = 2
        RelayController     = 3
        GyroSensor          = 4
        AccelerometerSensor = 5
        UltrasonicSensor    = 6
        GearToothSensor     = 7
        FirmwareUpdate      = 31

    class Manufacturer:
        BroadcastMessages   = 0
        NationalInstruments = 1
        TexasInstruments    = 2
        DEKA                = 3

    class APIClass:
        VoltageControl      = 0
        SpeedControl        = 1
        VoltageCompControl  = 2
        PositionControl     = 3
        CurrentControl      = 4
        Status              = 5
        Configuration       = 7
        Acknowledge         = 8

    class SystemControl:
        SystemHalt          = 0
        SystemReset         = 1
        DeviceAssignment    = 2
        DeviceQuery         = 3
        Heartbeat           = 5
        SynchronousUpdate   = 6
        FirmwareUpdate      = 7
        FirmwareVersion     = 8
        Enumeration         = 9
        SystemResume        = 10

    class SpeedControl:
        SpeedModeEnable     = 0
        SpeedModeDisable    = 1
        SpeedSet            = 2
        SpeedPConstant      = 3
        SpeedIConstant      = 4
        SpeedDConstant      = 5
        SpeedReference      = 6

    class MotorControlStatus:
        OutputVoltagePerc   = 0
        BusVoltage          = 1
        Current             = 2
        Temperature         = 3
        Position            = 4
        Speed               = 5
        Limit               = 6
        Fault               = 7
        Power               = 8
        ControlMode         = 9
        OutputVoltageVolts  = 10

    class MotorControlConfiguration:
        NumBrushes          = 0
        NumEncoderLines     = 1
        NumPotTurns         = 2
        BreakCoastSetting   = 3
        LimitMode           = 4
        ForwardDirLimit     = 5
        ReverseDirLimit     = 6
        MaximumOutputVolt   = 7
        FaultTime           = 8

    class SpeedReference:
        Positive   = 0
        Negative   = 2
        Quadrature = 3

    enumerate_timeout = timeval(milliseconds=80)

    def __init__(self, bridge):
        self.bridge = bridge

    def halt(self):
        '''
        Stop driving all motors and go into a neutral state. No motors can not
        be driven again until either a resume has been issued.
        '''
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.BroadcastMessages,
            manufacturer  = self.Manufacturer.BroadcastMessages,
            api_class     = 0, # arbitrary
            api_key       = self.SystemControl.Halt,
            device_number = 0, # arbitrary
            payload       = b''
        ))

    def resume(self):
        '''
        Enable motor control. See halt() for more information.
        '''
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.BroadcastMessages,
            manufacturer  = self.Manufacturer.BroadcastMessages,
            api_class     = 0, # arbitrary
            api_key       = self.SystemControl.Resume,
            device_number = 0, # arbitrary
            payload       = b''
        ))

    def update(self, mask):
        '''
        Trigger a synchronous update all motor controllers with a group that
        matches the bitmask.
        '''
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.BroadcastMessages,
            manufacturer  = self.Manufacturer.BroadcastMessages,
            api_class     = 0, # arbitrary
            api_key       = self.SystemControl.SynchronousUpdate,
            device_number = 0, # arbitrary
            payload       = struct.pack('B', mask)
        ))

    def heartbeat(self):
        '''
        Reset the watchdogs on all connected motor controllers. Every motor
        controller must receive a minimum of one CAN message every 100 ms to
        avoid entering a fault state.
        '''
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.BroadcastMessages,
            manufacturer  = self.Manufacturer.BroadcastMessages,
            api_class     = 0, # arbitrary
            api_key       = self.SystemControl.Heartbeat,
            device_number = 0, # arbitrary
            payload       = b''
        ))

    def enumerate(self):
        '''
        Obtain a list of motor controllers connected to this CAN bus bridge by
        requesting each controller to enumerate itself. This enumeration takes
        a minimum of 80 ms.
        '''
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.BroadcastMessages,
            manufacturer  = self.Manufacturer.BroadcastMessages,
            api_class     = 0, # arbitrary
            api_key       = self.SystemControl.Enumeration,
            device_number = 0, # arbitrary
            payload       = b''
        ))

        controllers = list()
        end = datetime.now() + self.enumerate_timeout
        while datetime.now() < end:
            remaining = (end - datetime.now()).total_seconds()
            msg = self.bridge.recv_message(timeout=remaining)

            if msg:
                ctrl = MotorCtrl(
                    device_type=msg.device_type,
                    manufacturer=msg.manufacturer,
                    device_number=msg.device_number
                )
                controllers.append(ctrl)

        return ids

class Jaguar:
    def __init__(self, bridge, device_number):
        self.bridge = bridge
        self.number = device_number

        # Use a quadrature encoder to measure velocity.
        payload = struct.pack('B', JaguarCAN.SpeedReference.Quadrature)
        self._send_message(JaguarCAN.SpeedControl.SpeedReference, payload)

    def enable(self):
        '''
        Enables speed control and sets the output to neutral.
        '''
        self._send_message(JaguarCAN.SpeedControl.SpeedModeEnable, b'')

    def disable(self):
        '''
        Disables speed control, returns to the default control mode, and sets
        the output to neutral.
        '''
        self._send_message(JaguarCAN.SpeedControl.SpeedModeDisable, b'')

    def set_speed(self, speed, group=0):
        '''
        Sets the target rotational speed in revolutions per minute.
        '''
        payload_speed = struct.pack('<i', ???)
        payload_group = struct.pack('B', group)
        payload       = payload_speed + payload_group
        self._send_message(self.SpeedControl.SpeedSet, payload)

    def set_constants(self, kp, ki, kd):
        payload_kp = struct.pack('<i', ???)
        payload_ki = struct.pack('<i', ???)
        payload_kd = struct.pack('<i', ???)
        self._send_message(self.SpeedControl.SetPConstant, payload_kp)
        self._send_message(self.SpeedControl.SetIConstant, payload_ki)
        self._send_message(self.SpeedControl.SetDConstant, payload_kd)

    def _send_message(self, api_key, payload):
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.MotorController,
            manufacturer  = self.Manufacturer.TexasInstruments,
            api_class     = self.APIClass.SpeedControl,
            api_key       = api_key,
            device_number = self.number,
            payload       = payload
        ))

def main():
    serial = serial.Serial(
        baudrate = 115200,
        bytesize = serial.EIGHTBITS,
        parity   = serial.PARITY_NONE,
        stopbits = serial.STOPBITS_ONE,
        timeout  = None
    )
    packetizer = Packetizer(0xff, 0xfe, 0xfe, 0xfd)
    jaguar = JaguarUART(serial, packetizer)

if __name__ == '__main__':
    main()
<commit_msg>Factored broadcast code into a method.<commit_after>import bitstring, serial, struct
from collections import namedtuple, OrderedDict
from datetime import datetime, timedelta

SerialPayload = namedtuple('SerialPayload', [ 'device_id', 'payload' ])
GenericCANMsg = namedtuple('GenericCANMsg', [ 'device_type', 'manufacturer',
                                              'api_class', 'api_key',
                                              'device_number', 'payload' ])
MotorCtrl     = namedtuple('MotorCtrl',     [ 'device_type', 'manufacturer',
                                              'device_number' ])

class Packetizer:
    def __init__(self, sof, esc, sof_esc, esc_esc):
        self.sof = bytearray([ sof ])
        self.esc = bytearray([ esc ])
        self.esc_sof = bytearray([ sof_esc ])
        self.esc_esc = bytearray([ esc_esc ])
        self._reset()

    def recv_byte(self, byte):
        byte_dec = self._decode(self.last, byte)

        # Every start of frame (SOF) byte starts a new frame because it is
        # otherwise escaped.
        if byte == self.sof:
            self._reset()
        # Next byte is the total number of bytes in the packet. Note that this
        # could be encoded if it equals 255, although this should never occur
        # in practice when framing CANbus messages.
        elif self.offset == 1:
            self.count   = ord(byte_dec)
            self.offset += 1
        elif byte_dec:
            self.packet.extend(byte_dec)
            self.offset += 1

        self.last = byte

        # Done receiving the entire packet, so add it to the queue. Padding
        # compensates for fields in the packet (e.g. device id) that don't
        # contribute to the byte count.
        if self.valid and self.offset == self.count + 2:
            new_packet = self.packet
            self._reset()
            return new_packet
        else:
            return None

    def frame_bytes(self, payload):
        count = len(payload)
        contents = struct.pack('B{}s'.format(count), count, payload)
        return self.sof + self._encode(contents)

    def _reset(self):
        self.offset = 1
        self.count  = 0
        self.last   = None
        self.packet = bytearray()
        self.valid  = True

    def _error(self):
        self.valid = False

    def _decode(self, last, curr):
        if last == self.esc:
            if curr == self.esc_esc:
                return self.esc
            elif curr == self.esc_sof:
                return self.sof
        elif curr == self.esc:
            return None
        else:
            return curr

    def _encode(self, raw):
        encoded = bytearray()
        for curr in raw:
            if curr == self.sof:
                encoded.extend(self.esc)
                encoded.extend(self.esc_sof)
            elif curr == self.esc:
                encoded.extend(self.esc)
                encoded.extend(self.esc_esc)
            else:
                encoded.extend(curr)
        return encoded

class JaguarUART:
    header_fields = OrderedDict([
        ('resv',          'uint:3'),
        ('device_type',   'uint:5'),
        ('manufacturer',  'uint:8'),
        ('api_class',     'uint:6'),
        ('api_key',       'uint:4'),
        ('device_number', 'uint:6')
    ])
    header_fmt = [ field + '=' + name for name, field in header_fields.items() ]

    def __init__(self, serial, packetizer):
        self.serial = serial
        self.packetizer = packetizer
        self.alive = True

    def close(self):
        self.alive = False

    @classmethod
    def parse_message(cls, packet):
        header_raw = bitstring.BitArray(bytes=packet[0:4])
        header_raw.byteswap()

        header_list = header_raw.unpack(cls.header_fmt)
        header = dict(zip(cls.header_fields.keys(), header_list))
        del header['resv']
        return GenericCANMsg(payload=packet[4:], **header)

    @classmethod
    def generate_message(cls, msg):
        header = bitstring.pack(cls.header_fmt, resv=0, **msg._asdict())
        header.byteswap()
        return header.bytes + msg.payload

    def send_message(self, msg):
        packed = self.generate_message(msg)
        framed = self.packetizer.frame_bytes(packed)
        self.serial.write(framed)

    def recv_message(self, timeout=None):
        self.timeout = timeout
        while True:
            # This may timeout and return zero bytes, so timeout here as well.
            byte = self.serial.read(1)
            if byte:
                return None

            # Feed the packetizing state machine the next character.
            packet = self.packetizer.recv_byte(byte)
            if packet:
                return self.parse_message(packet)

    def pending_data(self):
        return self.serial.inWaiting() > 0

class JaguarCAN:
    class DeviceType:
        BroadcastMessages   = 0
        RobotController     = 1
        MotorController     = 2
        RelayController     = 3
        GyroSensor          = 4
        AccelerometerSensor = 5
        UltrasonicSensor    = 6
        GearToothSensor     = 7
        FirmwareUpdate      = 31

    class Manufacturer:
        BroadcastMessages   = 0
        NationalInstruments = 1
        TexasInstruments    = 2
        DEKA                = 3

    class APIClass:
        VoltageControl      = 0
        SpeedControl        = 1
        VoltageCompControl  = 2
        PositionControl     = 3
        CurrentControl      = 4
        Status              = 5
        Configuration       = 7
        Acknowledge         = 8

    class SystemControl:
        SystemHalt          = 0
        SystemReset         = 1
        DeviceAssignment    = 2
        DeviceQuery         = 3
        Heartbeat           = 5
        SynchronousUpdate   = 6
        FirmwareUpdate      = 7
        FirmwareVersion     = 8
        Enumeration         = 9
        SystemResume        = 10

    class SpeedControl:
        SpeedModeEnable     = 0
        SpeedModeDisable    = 1
        SpeedSet            = 2
        SpeedPConstant      = 3
        SpeedIConstant      = 4
        SpeedDConstant      = 5
        SpeedReference      = 6

    class MotorControlStatus:
        OutputVoltagePerc   = 0
        BusVoltage          = 1
        Current             = 2
        Temperature         = 3
        Position            = 4
        Speed               = 5
        Limit               = 6
        Fault               = 7
        Power               = 8
        ControlMode         = 9
        OutputVoltageVolts  = 10

    class MotorControlConfiguration:
        NumBrushes          = 0
        NumEncoderLines     = 1
        NumPotTurns         = 2
        BreakCoastSetting   = 3
        LimitMode           = 4
        ForwardDirLimit     = 5
        ReverseDirLimit     = 6
        MaximumOutputVolt   = 7
        FaultTime           = 8

    class SpeedReference:
        Positive   = 0
        Negative   = 2
        Quadrature = 3

    enumerate_timeout = timeval(milliseconds=80)

    def __init__(self, bridge):
        self.bridge = bridge

    def _broadcast(self, api_key, payload):
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.BroadcastMessages,
            manufacturer  = self.Manufacturer.BroadcastMessages,
            api_class     = 0, # arbitrary
            api_key       = api_key,
            device_number = 0, # arbitrary
            payload       = payload
        ))

    def halt(self):
        '''
        Stop driving all motors and go into a neutral state. No motors can not
        be driven again until either a resume has been issued.
        '''
        self._broadcast(self.SystemControl.Halt, b'')

    def resume(self):
        '''
        Enable motor control. See halt() for more information.
        '''
        self._broadcast(self.SystemControl.Resume, b'')

    def update(self, mask):
        '''
        Trigger a synchronous update all motor controllers with a group that
        matches the bitmask.
        '''
        self._broadcast(self.SynchronousUpdate, struct.pack('B', mask))

    def heartbeat(self):
        '''
        Reset the watchdogs on all connected motor controllers. Every motor
        controller must receive a minimum of one CAN message every 100 ms to
        avoid entering a fault state.
        '''
        self._broadcast(self.Heartbeat, b'')

    def enumerate(self):
        '''
        Obtain a list of motor controllers connected to this CAN bus bridge by
        requesting each controller to enumerate itself. This enumeration takes
        a minimum of 80 ms.
        '''
        self._broadcast(self.Enumeration, b'')

        controllers = list()
        end = datetime.now() + self.enumerate_timeout
        while datetime.now() < end:
            remaining = (end - datetime.now()).total_seconds()
            msg = self.bridge.recv_message(timeout=remaining)

            if msg:
                ctrl = MotorCtrl(
                    device_type=msg.device_type,
                    manufacturer=msg.manufacturer,
                    device_number=msg.device_number
                )
                controllers.append(ctrl)

        return ids

class Jaguar:
    def __init__(self, bridge, device_number):
        self.bridge = bridge
        self.number = device_number

        # Use a quadrature encoder to measure velocity.
        payload = struct.pack('B', JaguarCAN.SpeedReference.Quadrature)
        self._send_message(JaguarCAN.SpeedControl.SpeedReference, payload)

    def enable(self):
        '''
        Enables speed control and sets the output to neutral.
        '''
        self._send_message(JaguarCAN.SpeedControl.SpeedModeEnable, b'')

    def disable(self):
        '''
        Disables speed control, returns to the default control mode, and sets
        the output to neutral.
        '''
        self._send_message(JaguarCAN.SpeedControl.SpeedModeDisable, b'')

    def set_speed(self, speed, group=0):
        '''
        Sets the target rotational speed in revolutions per minute.
        '''
        payload_speed = struct.pack('<i', ???)
        payload_group = struct.pack('B', group)
        payload       = payload_speed + payload_group
        self._send_message(self.SpeedControl.SpeedSet, payload)

    def set_constants(self, kp, ki, kd):
        payload_kp = struct.pack('<i', ???)
        payload_ki = struct.pack('<i', ???)
        payload_kd = struct.pack('<i', ???)
        self._send_message(self.SpeedControl.SetPConstant, payload_kp)
        self._send_message(self.SpeedControl.SetIConstant, payload_ki)
        self._send_message(self.SpeedControl.SetDConstant, payload_kd)

    def _send_message(self, api_key, payload):
        self.bridge.send_message(GenericCANMsg(
            device_type   = self.DeviceType.MotorController,
            manufacturer  = self.Manufacturer.TexasInstruments,
            api_class     = self.APIClass.SpeedControl,
            api_key       = api_key,
            device_number = self.number,
            payload       = payload
        ))

def main():
    serial = serial.Serial(
        baudrate = 115200,
        bytesize = serial.EIGHTBITS,
        parity   = serial.PARITY_NONE,
        stopbits = serial.STOPBITS_ONE,
        timeout  = None
    )
    packetizer = Packetizer(0xff, 0xfe, 0xfe, 0xfd)
    jaguar = JaguarUART(serial, packetizer)

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


<commit_before># -*- coding: utf-8 -*-
# Copyright: Damien Elmes <anki@ichi2.net>
# License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html

import re
import os
import shutil
import zipfile

from anki import Collection
from anki.hooks import runHook
from anki.lang import _
from anki.utils import ids2str, json, splitFields

class Exporter(object):
    def __init__(self, col, did=None):
        self.col = col
        self.did = did

    def exportInto(self, path):
        self._escapeCount = 0
        file = open(path, "wb")
        self.doExport(file)
        file.close()

    def escapeText(self, text):
        "Escape newlines, tabs and CSS."
        text = text.replace("\n", "
")
        text = text.replace("\t", " " * 8)
        text = re.sub("(?i)", "", text)
        return text

    def cardIds(self):
        if not self.did:
            cids = self.col.db.list("select id from cards")
        else:
            cids = self.col.decks.cids(self.did, children=True)
        self.count = len(cids)
        return cids

# Cards as TSV
######################################################################

class TextCardExporter(Exporter):

    key = _("Cards in Plain Text")
    ext = ".txt"
    hideTags = True

    def __init__(self, col):
        Exporter.__init__(self, col)

    def doExport(self, file):
        ids = sorted(self.cardIds())
        strids = ids2str(ids)

        def esc(s):
            # strip off the repeated question in answer if exists
            s = re.sub("(?si)^.*
\n*", "", s)
            return self.escapeText(s)
        out = ""
        for cid in ids:
            c = self.col.getCard(cid)
            out += esc(c.q())
            out += "\t" + esc(c.a()) + "\n"
        file.write(out.encode("utf-8"))

# Notes as TSV
######################################################################

class TextNoteExporter(Exporter):

    key = _("Notes in Plain Text")
    ext = ".txt"

    def __init__(self, col):
        Exporter.__init__(self, col)
        self.includeID = False
        self.includeTags = True

    def doExport(self, file):
        cardIds = self.cardIds()
        data = []
        for id, flds, tags in self.col.db.execute("""
select guid, flds, tags from notes
where id in
(select nid from cards
where cards.id in %s)""" % ids2str(cardIds)):
            row = []
            # note id
            if self.includeID:
                row.append(str(id))
            # fields
            row.extend([self.escapeText(f) for f in splitFields(flds)])
            # tags
            if self.includeTags:
                row.append(tags.strip())
            data.append("\t".join(row))
        self.count = len(data)
        out = "\n".join(data)
        file.write(out.encode("utf-8"))

# Anki decks
######################################################################
# media files are stored in self.mediaFiles, but not exported.

class AnkiExporter(Exporter):

    key = _("Anki 2.0 Deck")
    ext = ".anki2"

    def __init__(self, col):
        Exporter.__init__(self, col)
        self.includeSched = False
        self.includeMedia = True

    def exportInto(self, path):
        # create a new collection at the target
        try:
            os.unlink(path)
        except (IOError, OSError):
            pass
        self.dst = Collection(path)
        self.src = self.col
        # find cards
        if not self.did:
            cids = self.src.db.list("select id from cards")
        else:
            cids = self.src.decks.cids(self.did, children=True)
        # copy cards, noting used nids
        nids = {}
        data = []
        for row in self.src.db.execute(
                "select * from cards where id in " + ids2str(cids)):
            nids[row[1]] = True
            data.append(row)
        self.dst.db.executemany(
            "insert into cards values (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
            data)
        # notes
        strnids = ids2str(nids.keys())
        notedata = self.src.db.all("select * from notes where id in " +
                                   strnids)
        self.dst.db.executemany(
            "insert into notes values (?,?,?,?,?,?,?,?,?,?,?)",
            notedata)
        # models used by the notes
        mids = self.dst.db.list("select distinct mid from notes where id in " +
                                strnids)
        # card history and revlog
        if self.includeSched:
            data = self.src.db.all(
                "select * from revlog where cid in " + ids2str(cids))
            self.dst.db.executemany(
                "insert into revlog values (?,?,?,?,?,?,?,?,?)",
                data)
        else:
            # need to reset card state
            self.dst.sched.resetCards(cids)
        # models
        for m in self.src.models.all():
            if int(m['id']) in mids:
                self.dst.models.update(m)
        # decks
        if not self.did:
            dids = []
        else:
            dids = [self.did] + [
                x[1] for x in self.src.decks.children(self.did)]
        dconfs = {}
        for d in self.src.decks.all():
            if str(d['id']) == "1":
                continue
            if dids and d['id'] not in dids:
                continue
            if not d['dyn'] and d['conf'] != 1:
                if self.includeSched:
                    dconfs[d['conf']] = True
            if not self.includeSched:
                # scheduling not included, so reset deck settings to default
                d = dict(d)
                d['conf'] = 1
            self.dst.decks.update(d)
        # copy used deck confs
        for dc in self.src.decks.allConf():
            if dc['id'] in dconfs:
                self.dst.decks.updateConf(dc)
        # find used media
        media = {}
        self.mediaDir = self.src.media.dir()
        if self.includeMedia:
            for row in notedata:
                flds = row[6]
                mid = row[2]
                for file in self.src.media.filesInStr(mid, flds):
                    media[file] = True
            if self.mediaDir:
                for fname in os.listdir(self.mediaDir):
                    if fname.startswith("_"):
                        media[fname] = True
        self.mediaFiles = media.keys()
        self.dst.crt = self.src.crt
        # todo: tags?
        self.count = self.dst.cardCount()
        self.dst.setMod()
        self.postExport()
        self.dst.close()

    def postExport(self):
        # overwrite to apply customizations to the deck before it's closed,
        # such as update the deck description
        pass

# Packaged Anki decks
######################################################################

class AnkiPackageExporter(AnkiExporter):

    key = _("Anki Deck Package")
    ext = ".apkg"

    def __init__(self, col):
        AnkiExporter.__init__(self, col)

    def exportInto(self, path):
        # open a zip file
        z = zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED)
        # if all decks and scheduling included, full export
        if self.includeSched and not self.did:
            media = self.exportVerbatim(z)
        else:
            # otherwise, filter
            media = self.exportFiltered(z, path)
        # media map
        z.writestr("media", json.dumps(media))
        z.close()

    def exportFiltered(self, z, path):
        # export into the anki2 file
        colfile = path.replace(".apkg", ".anki2")
        AnkiExporter.exportInto(self, colfile)
        z.write(colfile, "collection.anki2")
        # and media
        self.prepareMedia()
        media = {}
        for c, file in enumerate(self.mediaFiles):
            c = str(c)
            mpath = os.path.join(self.mediaDir, file)
            if os.path.exists(mpath):
                z.write(mpath, c)
                media[c] = file
        # tidy up intermediate files
        os.unlink(colfile)
        os.unlink(path.replace(".apkg", ".media.db"))
        shutil.rmtree(path.replace(".apkg", ".media"))
        return media

    def exportVerbatim(self, z):
        # close our deck & write it into the zip file, and reopen
        self.count = self.col.cardCount()
        self.col.close()
        z.write(self.col.path, "collection.anki2")
        self.col.reopen()
        # copy all media
        if not self.includeMedia:
            return {}
        media = {}
        mdir = self.col.media.dir()
        for c, file in enumerate(os.listdir(mdir)):
            c = str(c)
            mpath = os.path.join(mdir, file)
            if os.path.exists(mpath):
                z.write(mpath, c)
                media[c] = file
        return media

    def prepareMedia(self):
        # chance to move each file in self.mediaFiles into place before media
        # is zipped up
        pass

# Export modules
##########################################################################

def exporters():
    def id(obj):
        return ("%s (*%s)" % (obj.key, obj.ext), obj)
    exps = [
        id(AnkiPackageExporter),
        id(TextNoteExporter),
        id(TextCardExporter),
    ]
    runHook("exportersList", exps)
    return exps
<commit_msg>Take out unused variable, per flake8 warning.<commit_after># -*- coding: utf-8 -*-
# Copyright: Damien Elmes <anki@ichi2.net>
# License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html

import re
import os
import shutil
import zipfile

from anki import Collection
from anki.hooks import runHook
from anki.lang import _
from anki.utils import ids2str, json, splitFields

class Exporter(object):
    def __init__(self, col, did=None):
        self.col = col
        self.did = did

    def exportInto(self, path):
        self._escapeCount = 0
        file = open(path, "wb")
        self.doExport(file)
        file.close()

    def escapeText(self, text):
        "Escape newlines, tabs and CSS."
        text = text.replace("\n", "
")
        text = text.replace("\t", " " * 8)
        text = re.sub("(?i)", "", text)
        return text

    def cardIds(self):
        if not self.did:
            cids = self.col.db.list("select id from cards")
        else:
            cids = self.col.decks.cids(self.did, children=True)
        self.count = len(cids)
        return cids

# Cards as TSV
######################################################################

class TextCardExporter(Exporter):

    key = _("Cards in Plain Text")
    ext = ".txt"
    hideTags = True

    def __init__(self, col):
        Exporter.__init__(self, col)

    def doExport(self, file):
        ids = sorted(self.cardIds())
        # The ids2str has no side effect, so take out redundant call
        # strids = ids2str(ids)

        def esc(s):
            # strip off the repeated question in answer if exists
            s = re.sub("(?si)^.*
\n*", "", s)
            return self.escapeText(s)
        out = ""
        for cid in ids:
            c = self.col.getCard(cid)
            out += esc(c.q())
            out += "\t" + esc(c.a()) + "\n"
        file.write(out.encode("utf-8"))

# Notes as TSV
######################################################################

class TextNoteExporter(Exporter):

    key = _("Notes in Plain Text")
    ext = ".txt"

    def __init__(self, col):
        Exporter.__init__(self, col)
        self.includeID = False
        self.includeTags = True

    def doExport(self, file):
        cardIds = self.cardIds()
        data = []
        for id, flds, tags in self.col.db.execute("""
select guid, flds, tags from notes
where id in
(select nid from cards
where cards.id in %s)""" % ids2str(cardIds)):
            row = []
            # note id
            if self.includeID:
                row.append(str(id))
            # fields
            row.extend([self.escapeText(f) for f in splitFields(flds)])
            # tags
            if self.includeTags:
                row.append(tags.strip())
            data.append("\t".join(row))
        self.count = len(data)
        out = "\n".join(data)
        file.write(out.encode("utf-8"))

# Anki decks
######################################################################
# media files are stored in self.mediaFiles, but not exported.

class AnkiExporter(Exporter):

    key = _("Anki 2.0 Deck")
    ext = ".anki2"

    def __init__(self, col):
        Exporter.__init__(self, col)
        self.includeSched = False
        self.includeMedia = True

    def exportInto(self, path):
        # create a new collection at the target
        try:
            os.unlink(path)
        except (IOError, OSError):
            pass
        self.dst = Collection(path)
        self.src = self.col
        # find cards
        if not self.did:
            cids = self.src.db.list("select id from cards")
        else:
            cids = self.src.decks.cids(self.did, children=True)
        # copy cards, noting used nids
        nids = {}
        data = []
        for row in self.src.db.execute(
                "select * from cards where id in " + ids2str(cids)):
            nids[row[1]] = True
            data.append(row)
        self.dst.db.executemany(
            "insert into cards values (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
            data)
        # notes
        strnids = ids2str(nids.keys())
        notedata = self.src.db.all("select * from notes where id in " +
                                   strnids)
        self.dst.db.executemany(
            "insert into notes values (?,?,?,?,?,?,?,?,?,?,?)",
            notedata)
        # models used by the notes
        mids = self.dst.db.list("select distinct mid from notes where id in " +
                                strnids)
        # card history and revlog
        if self.includeSched:
            data = self.src.db.all(
                "select * from revlog where cid in " + ids2str(cids))
            self.dst.db.executemany(
                "insert into revlog values (?,?,?,?,?,?,?,?,?)",
                data)
        else:
            # need to reset card state
            self.dst.sched.resetCards(cids)
        # models
        for m in self.src.models.all():
            if int(m['id']) in mids:
                self.dst.models.update(m)
        # decks
        if not self.did:
            dids = []
        else:
            dids = [self.did] + [
                x[1] for x in self.src.decks.children(self.did)]
        dconfs = {}
        for d in self.src.decks.all():
            if str(d['id']) == "1":
                continue
            if dids and d['id'] not in dids:
                continue
            if not d['dyn'] and d['conf'] != 1:
                if self.includeSched:
                    dconfs[d['conf']] = True
            if not self.includeSched:
                # scheduling not included, so reset deck settings to default
                d = dict(d)
                d['conf'] = 1
            self.dst.decks.update(d)
        # copy used deck confs
        for dc in self.src.decks.allConf():
            if dc['id'] in dconfs:
                self.dst.decks.updateConf(dc)
        # find used media
        media = {}
        self.mediaDir = self.src.media.dir()
        if self.includeMedia:
            for row in notedata:
                flds = row[6]
                mid = row[2]
                for file in self.src.media.filesInStr(mid, flds):
                    media[file] = True
            if self.mediaDir:
                for fname in os.listdir(self.mediaDir):
                    if fname.startswith("_"):
                        media[fname] = True
        self.mediaFiles = media.keys()
        self.dst.crt = self.src.crt
        # todo: tags?
        self.count = self.dst.cardCount()
        self.dst.setMod()
        self.postExport()
        self.dst.close()

    def postExport(self):
        # overwrite to apply customizations to the deck before it's closed,
        # such as update the deck description
        pass

# Packaged Anki decks
######################################################################

class AnkiPackageExporter(AnkiExporter):

    key = _("Anki Deck Package")
    ext = ".apkg"

    def __init__(self, col):
        AnkiExporter.__init__(self, col)

    def exportInto(self, path):
        # open a zip file
        z = zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED)
        # if all decks and scheduling included, full export
        if self.includeSched and not self.did:
            media = self.exportVerbatim(z)
        else:
            # otherwise, filter
            media = self.exportFiltered(z, path)
        # media map
        z.writestr("media", json.dumps(media))
        z.close()

    def exportFiltered(self, z, path):
        # export into the anki2 file
        colfile = path.replace(".apkg", ".anki2")
        AnkiExporter.exportInto(self, colfile)
        z.write(colfile, "collection.anki2")
        # and media
        self.prepareMedia()
        media = {}
        for c, file in enumerate(self.mediaFiles):
            c = str(c)
            mpath = os.path.join(self.mediaDir, file)
            if os.path.exists(mpath):
                z.write(mpath, c)
                media[c] = file
        # tidy up intermediate files
        os.unlink(colfile)
        os.unlink(path.replace(".apkg", ".media.db"))
        shutil.rmtree(path.replace(".apkg", ".media"))
        return media

    def exportVerbatim(self, z):
        # close our deck & write it into the zip file, and reopen
        self.count = self.col.cardCount()
        self.col.close()
        z.write(self.col.path, "collection.anki2")
        self.col.reopen()
        # copy all media
        if not self.includeMedia:
            return {}
        media = {}
        mdir = self.col.media.dir()
        for c, file in enumerate(os.listdir(mdir)):
            c = str(c)
            mpath = os.path.join(mdir, file)
            if os.path.exists(mpath):
                z.write(mpath, c)
                media[c] = file
        return media

    def prepareMedia(self):
        # chance to move each file in self.mediaFiles into place before media
        # is zipped up
        pass

# Export modules
##########################################################################

def exporters():
    def id(obj):
        return ("%s (*%s)" % (obj.key, obj.ext), obj)
    exps = [
        id(AnkiPackageExporter),
        id(TextNoteExporter),
        id(TextCardExporter),
    ]
    runHook("exportersList", exps)
    return exps
<|endoftext|>