python: best way to verify conditions and raise Exception

hi i have to verify if a vector contains all 0 or 1 and if not raise exception:

def assert_all_zero_or_one(vector):
    if set(vector)=={0}: return 0
    if set(vector)=={1}: return 1
    raise TypeError

with this exceution

assert_all_zero_or_one([1,1,1]) # return 1
assert_all_zero_or_one([0,0]) # return 0
assert_all_zero_or_one([1,0,1]) # raise TypeError

i don't开发者_如何学Go like this solution.. there is a best way to do it with python?

I think your solution conveys your intent well. You could also do

def assert_all_zero_or_one(vector):
    if set(vector) not in ({0}, {1}): raise TypeError
    return vector[0]

so you build set(vector) only once, but I think yours is easier to understand.

How's this?

def assert_all_zero_or_one(vector):
    if all(v==1 for v in vector): return 1
    elif all(v==0 for v in vector): return 0
    else: raise TypeError

Reads quite nicely, I think.

If you prefer short and cryptic:

def assert_all_zero_or_one(vector):
    a, = set(vector)
    return a

Although that gives you a ValueError rather than a TypeError.

def allOneOf(items, ids):
    first = items[0]
    if first in ids and all(i==first for i in items):
        return first
    else:
        raise TypeError()

assert_all_zero_or_one = (lambda vector: allOneOf(vector, set([0,1])))

You can also do something like this.

import functools
import operator

def assert_all_zero_or_one(vector):
    if ((functools.reduce(operator.__mul__,vector) == 1) or (functools.reduce(operator.__add__,vector) == 0)):
            return vector[0]
    else:   
        raise TypeError("Your vector should be homogenous")

def assert_all_zero_or_one(vector):
    if len(vector) == sum(vector):
        return True
    elif sum(vector) == 0:
        return True
    return False

This should work nicely.