1# A minimal client for Mercurial's command server 2 3from __future__ import absolute_import, print_function 4 5import io 6import os 7import re 8import signal 9import socket 10import struct 11import subprocess 12import sys 13import time 14 15if sys.version_info[0] >= 3: 16 stdout = sys.stdout.buffer 17 stderr = sys.stderr.buffer 18 stringio = io.BytesIO 19 20 def bprint(*args): 21 # remove b'' as well for ease of test migration 22 pargs = [re.sub(br'''\bb(['"])''', br'\1', b'%s' % a) for a in args] 23 stdout.write(b' '.join(pargs) + b'\n') 24 25 26else: 27 import cStringIO 28 29 stdout = sys.stdout 30 stderr = sys.stderr 31 stringio = cStringIO.StringIO 32 bprint = print 33 34 35def connectpipe(path=None, extraargs=()): 36 cmdline = [b'hg', b'serve', b'--cmdserver', b'pipe'] 37 if path: 38 cmdline += [b'-R', path] 39 cmdline.extend(extraargs) 40 41 def tonative(cmdline): 42 if os.name != 'nt': 43 return cmdline 44 return [arg.decode("utf-8") for arg in cmdline] 45 46 server = subprocess.Popen( 47 tonative(cmdline), stdin=subprocess.PIPE, stdout=subprocess.PIPE 48 ) 49 50 return server 51 52 53class unixconnection(object): 54 def __init__(self, sockpath): 55 self.sock = sock = socket.socket(socket.AF_UNIX) 56 sock.connect(sockpath) 57 self.stdin = sock.makefile('wb') 58 self.stdout = sock.makefile('rb') 59 60 def wait(self): 61 self.stdin.close() 62 self.stdout.close() 63 self.sock.close() 64 65 66class unixserver(object): 67 def __init__(self, sockpath, logpath=None, repopath=None): 68 self.sockpath = sockpath 69 cmdline = [b'hg', b'serve', b'--cmdserver', b'unix', b'-a', sockpath] 70 if repopath: 71 cmdline += [b'-R', repopath] 72 if logpath: 73 stdout = open(logpath, 'a') 74 stderr = subprocess.STDOUT 75 else: 76 stdout = stderr = None 77 self.server = subprocess.Popen(cmdline, stdout=stdout, stderr=stderr) 78 # wait for listen() 79 while self.server.poll() is None: 80 if os.path.exists(sockpath): 81 break 82 time.sleep(0.1) 83 84 def connect(self): 85 return unixconnection(self.sockpath) 86 87 def shutdown(self): 88 os.kill(self.server.pid, signal.SIGTERM) 89 self.server.wait() 90 91 92def writeblock(server, data): 93 server.stdin.write(struct.pack(b'>I', len(data))) 94 server.stdin.write(data) 95 server.stdin.flush() 96 97 98def readchannel(server): 99 data = server.stdout.read(5) 100 if not data: 101 raise EOFError 102 channel, length = struct.unpack('>cI', data) 103 if channel in b'IL': 104 return channel, length 105 else: 106 return channel, server.stdout.read(length) 107 108 109def sep(text): 110 return text.replace(b'\\', b'/') 111 112 113def runcommand( 114 server, args, output=stdout, error=stderr, input=None, outfilter=lambda x: x 115): 116 bprint(b'*** runcommand', b' '.join(args)) 117 stdout.flush() 118 server.stdin.write(b'runcommand\n') 119 writeblock(server, b'\0'.join(args)) 120 121 if not input: 122 input = stringio() 123 124 while True: 125 ch, data = readchannel(server) 126 if ch == b'o': 127 output.write(outfilter(data)) 128 output.flush() 129 elif ch == b'e': 130 error.write(data) 131 error.flush() 132 elif ch == b'I': 133 writeblock(server, input.read(data)) 134 elif ch == b'L': 135 writeblock(server, input.readline(data)) 136 elif ch == b'm': 137 bprint(b"message: %r" % data) 138 elif ch == b'r': 139 (ret,) = struct.unpack('>i', data) 140 if ret != 0: 141 bprint(b' [%d]' % ret) 142 return ret 143 else: 144 bprint(b"unexpected channel %c: %r" % (ch, data)) 145 if ch.isupper(): 146 return 147 148 149def check(func, connect=connectpipe): 150 stdout.flush() 151 server = connect() 152 try: 153 return func(server) 154 finally: 155 server.stdin.close() 156 server.wait() 157 158 159def checkwith(connect=connectpipe, **kwargs): 160 def wrap(func): 161 return check(func, lambda: connect(**kwargs)) 162 163 return wrap 164