git » pymisc » master » tree

[master] / tcpmp.py

"""
tcpmp - TCP Message Passing
Alberto Bertogli (albertito@blitiri.com.ar)
------------------------------------------

This module implements a bidirectional channel with the Channel class, which
allows you to send and receive python objects (as long as they're pickable)
between two hosts, using TCP.

You should instance in each host, and then call .server() or .client(),
according to your status. Then you can begin sending and receiving objects
with .send() and .recv().

If the connection gets broken (one of the processes closes it or dies), send()
and recv() will raise a BrokenChannel exception, defined here.

It should work under any Unix, Windows and Mac systems.
"""

import os
import socket
import errno
import pickle

default_port = 1642

class BrokenChannel (Exception):
	"Exception raised when the channel gets broken."
	pass

class SockFD (object):
	"""File descriptor wrapper for socket objects. Implement .write(),
	.read() and .readline(), which is what pickle needs to work."""

	def __init__(self, sock):
		self.sock = sock
		self.fd = os.fdopen(sock.fileno())

	def write(self, s):
		self.sock.send(s)

	def read(self, size = -1):
		if size == -1:
			return self.fd.read()
		return self.fd.read(size)

	def readline(self, size = -1):
		return self.fd.readline(size)

	def close(self):
		#self.sock.close()
		self.fd.close()

	def __del__(self):
		self.close()


class Channel (object):
	"TCP message passing."

	def __init__(self):
		self.sock = None
		self.fd = None

	def __del__(self):
		if self.fd:
			self.fd.close()

	def client(self, ip, port = default_port):
		"Tell the object we're the client."
		s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		s.connect((ip, port))
		self.fd = SockFD(s)

	def server(self, listen_ip = '0.0.0.0', port = default_port):
		"Tell the object we're the server."
		s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		s.bind((listen_ip, port))
		s.listen(1)
		conn, addr = s.accept()
		self.fd = SockFD(conn)

	def send(self, o):
		"Send an object."
		try:
			pickle.dump(o, self.fd, -1)
		except (socket.error, EOFError), info:
			self.fd.close()
			raise BrokenChannel, info

	def recv(self):
		"Receive an object."
		try:
			return pickle.load(self.fd)
		except (socket.error, EOFError), info:
			self.fd.close()
			raise BrokenChannel, info