# ---
# jupyter:
#   jupytext:
#     split_at_heading: true
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
#default_exp foundation
# -

#export
from fastcore.imports import *

from fastcore.test import *
from nbdev.showdoc import *

# # Core
#
# > Basic functions used in the fastai library

# export
defaults = SimpleNamespace()


# ## Metaclasses

# See this [blog post](https://realpython.com/python-metaclasses/) for more information about metaclasses. 
# - `PrePostInitMeta` ensures that the classes defined with it run `__pre_init__` and `__post_init__` (without having to write `self.__pre_init__()` and `self.__post_init__()`  in the actual `init`
# - `NewChkMeta` gives the `PrePostInitMeta` functionality and ensures classes defined with it don't re-create an object of their type whenever it's passed to the constructor
# - `BypassNewMeta` ensures classes defined with it can easily be casted form objects they subclass.

# **NB: While a class is being defined from a metaclass, the `cls` arg in `__new__` of the metaclass points to the metaclass itself while `base` points to the defined class.**
#
# **This behaviour changes as soon as the class has been defined. When it is now being instantiated, the `cls` in all the dunder methods of the metaclass now point to the class being instantaited instead of the metaclass **

#export 
def _rm_self(sig):
    # literally remove `self` from the class' __init__ signature

    # this takes the signature of a class' __init__, gets it's args/params dict, 
    # removes 'self' from the dict and then replaces the initial signature with the one without the 'self' 
    sigd = dict(sig.parameters)
    sigd.pop('self')
    return sig.replace(parameters=sigd.values())


inspect.signature(object.__init__)


class A(type):
    def __new__(cls, name, bases, dict):
        print('creating __new__')
        res = super().__new__(cls, name, bases, dict)
        print(res, cls, name, bases, dict)        
        print(res)

        return res


class B(metaclass = A): ... # the `cls` arg should point to the metaclass while `base` points to the defined class


# +
#export
# this is the metaclass that handles `self`-stripping classes that are instatiated by this meta
class FixSigMeta(type):
    "A metaclass that fixes the signature on classes that override __new__"
    def __new__(cls, name, bases, dict):
#         we want to change the __new__ from `type` in our new metaclass `FixSigMeta` 
#         to represent the base class that will be inheriting this meta
        res = super().__new__(cls, name, bases, dict)
        
#         if the init of the child class(base class) is not a slot wrapper(init of an object) ie the class has 
#         it's own init, strip that init signature of it's `self` arg by calling `_rm_self`
        if res.__init__ is not object.__init__: res.__signature__ = _rm_self(inspect.signature(res.__init__))
#         return this self-stripped base class
        return res


# -

class A(type):
    def __call__(cls, *args, **kwargs):
        print(cls)


class B(metaclass=A): ...


B(3) #it should point to the defined class

type(B.__new__(B))


#export
# this is the metaclass inheriting a metaclass that makes sure the optional `__pre_init__` and `__post_init__` are run if they exist
# We can use it to avoid ever having to call super() again in the init
class PrePostInitMeta(FixSigMeta):
    "A metaclass that calls optional `__pre_init__` and `__post_init__` methods"
    def __call__(cls, *args, **kwargs):
        #define a new subclass of the base class 
#         inotherwords
#         create a new definition of the base(defined) class inheriting this meta
        res = cls.__new__(cls)
#     if this subclass is of the same type with it's super class
        if type(res)==cls:
#             check if this subclass has `__pre_init__` if it does, call it
            if hasattr(res,'__pre_init__'): res.__pre_init__(*args,**kwargs)
#           now then run the init
            res.__init__(*args,**kwargs)
#           #if the class has a __post_init__, run it
            if hasattr(res,'__post_init__'): res.__post_init__(*args,**kwargs)
#         return the completely run class
        return res


show_doc(PrePostInitMeta, title_level=3)


# +
class _T(metaclass=PrePostInitMeta):
    def __pre_init__(self):  self.a  = 0; assert self.a==0
    def __init__(self,b=0):  self.a += 1; assert self.a==1
    def __post_init__(self): self.a += 1; assert self.a==2

t = _T()
test_eq(t.a, 2) #the post init if it exists is always the last to run


# -

#export
# NOTE: According to Jeremy, this is mostly used in L and fp16
# this is a meta that tries to make sure that if a subclass of a defined class is passed as an argument to the defined class, it knows and hence doesn't perform any operation
# it ensures if the class instance is passed as an arg to the class, the class does not change
class NewChkMeta(FixSigMeta):
    "Metaclass to avoid recreating object passed to constructor"
    def __call__(cls, x=None, *args, **kwargs):
#         if there's no args or kwarg input when the class is instantiated but x exists and it is an child of the defined class(ie it is an instance)
        if not args and not kwargs and x is not None and isinstance(x,cls):
#             break out of the and return the child (instance of defined class) as is
            x._newchk = 1
            return x
        res = super().__call__(*((x,) + args), **kwargs)

        res._newchk = 0
        return res


class _T(metaclass=NewChkMeta):
    "Testing"
    def __init__(self, o=None, b=1):
        self.foo = getattr(o,'foo',0) + 1
        self.b = b


# +
# building a full class with `type`

# R = type('R', (object, ), {'foo': 10})
def init(self, _a=0): #the _a is just there so that self will be recognized, doesn't change much unless another arg is added
    self.foo = _a
R = type('R', (_T, ), {'__init__': init}) #inorder to attach an init to a class defined with type, we have to use another func `init`
# -

issubclass(R, _T)

R().foo  #foo is an attribute of this which means that the `__init__` works!

R.__new__(_T) # how to create a new instance of a class

_T(R.__new__(_T))


# +
class _T2():
    def __init__(self, o): self.foo = getattr(o,'foo',0) + 1

t = _T(1)
test_eq(t.foo,1)
t2 = _T(t)
test_eq(t2.foo,1)
test_is(t,t2)
t3 = _T(t, b=2)
test_eq(t3.b, 2)
assert not t3 is t

t = _T2(1)
test_eq(t.foo,1)
t2 = _T2(t)
test_eq(t2.foo,2)

test_eq(_T.__doc__, "Testing")
test_eq(str(inspect.signature(_T)), '(o=None, b=1)')


# -

#export
# Not used anywhere else
class BypassNewMeta(FixSigMeta):
    "Metaclass: casts `x` to this class if it's of type `cls._bypass_type`, initializing with `_new_meta` if available"
    def __call__(cls, x=None, *args, **kwargs):
        if hasattr(cls, '_new_meta'): x = cls._new_meta(x, *args, **kwargs)
        elif not isinstance(x,getattr(cls,'_bypass_type',object)) or len(args) or len(kwargs):
            x = super().__call__(*((x,)+args), **kwargs)
        if cls!=x.__class__: x.__class__ = cls
        return x


# +
class T0: pass
class _T(T0, metaclass=BypassNewMeta):
    _bypass_type=T0
    def __init__(self,x): self.x=x

t = T0()
t.a = 1
t2 = _T(t)
test_eq(type(t2), _T)
test_eq(t2.a,1)
test_is(t2,t)
t = _T(2)
t.x = 2


# -

# ## Foundational functions

def f(mad='mad'):
    f = str(mad) + ' oh!'
    return f
f

# We can create a copy of this function simply by using `FunctionType(f.__code__, f.__globals__, f.__name__, f.__defaults__, f.__closure__)`

f_ = FunctionType(f.__code__, f.__globals__, f.__name__, f.__defaults__, f.__closure__)
f_

f(), f_()

copy_f = copy(f)
copy_f()

f.__dict__  #most funcs contain empty __dict__s by default


# +
#demonstrating MethodType
# we change a normal method to a class method
class A:
    def __init__(self): return None
    
# method defined outside class
def foo(self): 
    print('Hi')

print(MethodType(foo, A))

setattr(A, foo.__name__, MethodType(foo, A))

A.foo()   #it is a class method which can be called without instantiating the class


# -

#export
# func makes an exact copy of another func and also resets it `__qualname__`
def copy_func(f):
    "Copy a non-builtin function (NB `copy.copy` does not work for this)"
    if not isinstance(f,FunctionType): return copy(f)
    fn = FunctionType(f.__code__, f.__globals__, f.__name__, f.__defaults__, f.__closure__)
    fn.__dict__.update(f.__dict__)
    return fn


# **NB**: I can't quite see the need for `copy_func` because unlike what he says, copy.copy work perfectly well here.
#
# **EDIT**: Jeremy is the smart guy here stupid... The reason why Jeremy uses `copy_func` is because `copy.copy` doesn't copy the reset the __qualname__ of a copied func if it hs been changed. I have done a simulation of this difference in a cell below

# +
def func(b: [float, int], a): 
    return None

print(f'initial: {func.__qualname__}')
func.__qualname__ = '_T4.func'
print(f'changed qual: {func.__qualname__}')
print(' ')
func_b = copy(func)
func_b_ = copy_func(func)
print(f'using copy.copy: {func_b.__qualname__}') #copy doesn't reset it
print(f'using copy_func: {func_b_.__qualname__}') # copy func reset it


# -

# The trick to patching successfully is to chang the `__name__` of the func from `func_name` to the be `class_name.func_name` instead

# +
#export
# this decorator (notice the wrapper _inner) is used to add a new method to a predefined class or list of classes either 
# as a normal method classmethod or as a cls property
def patch_to(cls, as_prop=False, cls_method=False):
    "Decorator: add `f` to `cls`"
#     make sure it is a collection so we can loop over it
    if not isinstance(cls, (tuple,list)): cls=(cls,)
    def _inner(f):
        for c_ in cls:
#             create a copy of the func to be patched
            nf = copy_func(f)

            # `functools.update_wrapper` when passing patched function to `Pipeline`, so we do it manually
#             this is to update the info of the func `f` being wrapped so it won't be overridden by the wrapper
#             for o in functools.WRAPPER_ASSIGNMENTS: setattr(nf, o, getattr(f,o)) Jeremy style(from scratch method)
#             OR
            functools.update_wrapper(_inner, f) #Tendo style(pythonic style)
#           Note the reason why I'm not using a decorator @functools.wraps is because i don't have an _inner_inner 
#           func inside here. When patching we dont want the func to be called, we just want to pass it into the class as a method
            
#          This may be the most important bit, it sets the copy of the func being pathed `nf` as a method to the class `c_`4
#          Note: this isn't set to the main method because we may want to set the `f` as a class method or something else as shown below
            nf.__qualname__ = f"{c_.__name__}.{f.__name__}"
    
#             If we want this patched class to be a classmethod, We set the c.f.__name__ = method_bound(c_.nf)
            if cls_method:
                setattr(c_, f.__name__, MethodType(nf, c_))
            else:
                setattr(c_, f.__name__, property(nf) if as_prop else nf)
        return f
    return _inner


# +
class _T3(int): pass

@patch_to(_T3)
def func1(x, a): return x+a

@patch_to(_T3, cls_method=True)
def from_func2(cls, a): return cls(a)

t = _T3(1) #reason we can do this is because we are using int's __init__
test_eq(t.func1(2), 3)

t2 = _T3(2)
test_eq(t2, 2)

# +
_T3.from_func2(2) #A class method can be called without instantiating the method

#Tendo test
test_eq(_T3.from_func2(2), 2)


# -

# If `cls` is a tuple, `f` is added to all types in the tuple.

# +
class _T4(int): pass
@patch_to((_T3,_T4))
def func2(x, a): return x+2*a

t = _T3(1)
test_eq(t.func2(1), 3)
t = _T4(1)
test_eq(t.func2(1), 3)


# -

def h(a: [float, int], b:float): ...
print(h.__annotations__) # the annotations for every annotated func are stored in dunder annotations
next(iter(h.__annotations__.values())) # we only want to annotation of the very first arg


# This decorator uses `patch_to` to patch a func to a class which is based on the annotation of the 
# very first param of the func 
#export
def patch(f):
    "Decorator: add `f` to the first parameter's class (based on f's type annotations)"
    cls = next(iter(f.__annotations__.values())) # only the first param's/arg's type
    return patch_to(cls)(f)


# +
@patch
def func(x:_T3, a):
    "test"
    return x+a

t = _T3(1)
test_eq(t.func(3), 4)
test_eq(t.func.__qualname__, '_T3.func')


# -

# If annotation is a tuple, the function is added to all types in the tuple.

@patch
def func3(x:(_T3,_T4), a):
    "test"
    return x+2*a


# +
@patch
def func3(x:(_T3,_T4), a):
    "test"
    return x+2*a

t = _T3(1)
test_eq(t.func3(2), 5)
test_eq(t.func3.__qualname__, '_T3.func3')
t = _T4(1)
test_eq(t.func3(2), 5)
test_eq(t.func3.__qualname__, '_T4.func3')


# -

# decorator used to patch a method to a class based on the type of it's first arg as a class property 
#export
def patch_property(f):
    "Decorator: add `f` as a property to the first parameter's class (based on f's type annotations)"
    cls = next(iter(f.__annotations__.values()))
    return patch_to(cls, as_prop=True)(f)


# +
@patch_property
def prop(x:_T3): return x+1

t = _T3(1)
test_eq(t.prop, 2) # a class property is a method that behaves like an attribute in the class. If it has a setter it can be changed dynamically
# -

# how to make a func param/arg
inspect.Parameter('key', inspect.Parameter.KEYWORD_ONLY, default='value')


# this is used to make a func kwargs(key:value) param that has a default key-value of d(None) 
#export
def _mk_param(n,d=None): return inspect.Parameter(n, inspect.Parameter.KEYWORD_ONLY, default=d)


# representation of the params in a function
inspect.signature(f), str(inspect.signature(f))

# get the parameter dict or mapping proxy
inspect.signature(f).parameters, dict(inspect.signature(f).parameters), _mk_param('new_param', 'Hello')


# test to make sure func signatures match
def test_sig(f, b): test_eq(str(inspect.signature(f)), b)


# a decorator that allows us to use *kwargs while writing our function ans ensuring that autocomplete works becuase it saves 
# the signatures of the kwargs to the func's signature
#export
def use_kwargs_dict(keep=False, **kwargs):
    "Decorator: replace `**kwargs` in signature with `names` params"
    def _f(f):
#         obtain the Parameter signature
        sig = inspect.signature(f)
#         change this signature to a dict (mapping proxy) and then change it to a normal dict
        sigd = dict(sig.parameters)
#     remove the kwargs param from the dict
        k = sigd.pop('kwargs')
#     make new params using the key-value pairs passed into the decorator. Ensure that there is no key clash with the keys 
#     already existing in the func
        s2 = {n:_mk_param(n,d) for n,d in kwargs.items() if n not in sigd}
#     update the the signature dict of the function with these new Param key-value pairs
        sigd.update(s2)
        if keep: sigd['kwargs'] = k
#       replace the signature of the func
        f.__signature__ = sig.replace(parameters=sigd.values())
        return f
    return _f


@use_kwargs_dict(y=1,z=None)
def foo(a, b=1, **kwargs): pass
test_sig(foo, '(a, b=1, *, y=1, z=None)')


#this decorator is different from `use_kwargs_dict` because attaches non keyword args to the func signature dict
#export
def use_kwargs(names, keep=False):
    "Decorator: replace `**kwargs` in signature with `names` params"
    def _f(f):
        sig = inspect.signature(f)
        sigd = dict(sig.parameters)
        k = sigd.pop('kwargs')
        s2 = {n:_mk_param(n) for n in names if n not in sigd}  #<---this is the different bit. the default value `d` of None is used
        sigd.update(s2)
        if keep: sigd['kwargs'] = k
        f.__signature__ = sig.replace(parameters=sigd.values())
        return f
    return _f


# +
@use_kwargs(['y', 'z'])
def foo(a, b=1, **kwargs): pass
test_sig(foo, '(a, b=1, *, y=None, z=None)')

@use_kwargs(['y', 'z'], keep=True)
def foo(a, *args, b=1, **kwargs): pass
test_sig(foo, '(a, *args, b=1, y=None, z=None, **kwargs)')


# +
#extra test Tendo
@use_kwargs_dict(y_=1, z_=2)
@use_kwargs(['y', 'z'], keep=True)
def foo(a, b=1, **kwargs): pass

test_sig(foo, '(a, b=1, *, y=None, z=None, y_=1, z_=2)')

class Foo:
    def __init__(self): return None
    
    @classmethod
    @use_kwargs_dict(y=1, z=3)
    def say(self, b, c=9, **kwargs): None
        
test_sig(Foo.say, '(b, c=9, *, y=1, z=3)')


# -

# this decorator is used to transfer signature from one func `f` to another `to`. 
# This helps us with autocompletion when we a superclassing classes and using `**kwargs`
# It can be nested and you can also select the arg signatures to transfer
# export
def delegates(to=None, keep=False, but=None):
    "Decorator: replace `**kwargs` in signature with params from `to`"
    if but is None: but = []
    def _f(f):
#         if you're trying to delegate from one class to another. what you need to delegate is the classes' `__init__`
#        NOTE: in this case you don't pass arguments to the the class being delegated to
        if to is None: to_f,from_f = f.__base__.__init__,f.__init__
        else:          to_f,from_f = to,f
        from_f = getattr(from_f,'__func__',from_f)
        to_f = getattr(to_f,'__func__',to_f)
#         `__delwrap__` is a flag to prevent multiple nesting of delegates. If it occurs, break out
        if hasattr(from_f,'__delwrap__'): return f
        sig = inspect.signature(from_f)
        sigd = dict(sig.parameters)
        k = sigd.pop('kwargs')
#         ensure that the no value from the key-value kwargs are empty(not set) and that the key is not in `but` which 
#         should be excluded from delegation
        s2 = {k:v for k,v in inspect.signature(to_f).parameters.items()
              if v.default != inspect.Parameter.empty and k not in sigd and k not in but}
#       update the signatures
        sigd.update(s2)
#         this if else loop is present because inorder for us to use nested delegates, there has to be 
#         a **kwargs(ie the `keep` must be set to True)
        if keep: sigd['kwargs'] = k
#         if no keep don't allow nesting; set the flag
        else: from_f.__delwrap__ = to_f
#         replace them
        from_f.__signature__ = sig.replace(parameters=sigd.values())
#       since no copying was done from `f` to `from_f` we may as well return `from_f` here if we are not delegating inside a class
        return f
    return _f


# +
def basefoo(e, c=2): pass

@delegates(basefoo) # delegate all key:value pairs from basefoo to foo
def foo(a, b=1, **kwargs): pass
test_sig(foo, '(a, b=1, c=2)')

@delegates(basefoo, keep=True)
def foo(a, b=1, **kwargs): pass
test_sig(foo, '(a, b=1, c=2, **kwargs)')


@delegates(basefoo, but= ['c']) # delegate all args but c
def foo(a, b=1, **kwargs): pass
test_sig(foo, '(a, b=1)')

class _T():
    def foo(cls, a=1, b=2):
        pass
    
    @classmethod
    @delegates(foo)
    def bar(cls, c=3, **kwargs):
        pass

test_sig(_T.bar, '(c=3, a=1, b=2)')


# +
class BaseFoo:
    def __init__(self, c, d=2, e=3): pass  #note only kwargs are delegated

@delegates() #note that i have no args
class Foo(BaseFoo): # this super class we want to delagate from must be inherited
    def __init__(self, a, b=1, **kwargs): super().__init__(**kwargs)

test_sig(Foo, '(a, b=1, d=2, e=3)')


# +
#tendo test (nested delegation)
@delegates(BaseFoo.__init__)
@delegates(basefoo, keep=True) #the kwargs must be kept while nesting
def foo(a, b=1, **kwargs): pass

test_sig(foo, '(a, b=1, c=2, d=2, e=3)')
# -

MethodType(foo, 1)

print(isinstance(_T.foo,MethodType))
print(isinstance(_T.bar,MethodType))
print(isinstance(MethodType(foo, 1),MethodType))


# bind method foo to the class 1 (#numbers.Integral)
#export
def method(f):
    "Mark `f` as a method"
    # `1` is a dummy instance since Py3 doesn't allow `None` any more
    return MethodType(f, 1)


g = {'r': '3', 'w': '6'}
g.pop('f', '3'), g


#this decorator makes it possible for us to replace a method in a class with any other function so long as the 
# method is in the `cls._methods` list of the class `cls` and the `cls.__init__` has `kwargs`
#export
def _funcs_kwargs(cls, as_method):
    old_init = cls.__init__
    def _init(self, *args, **kwargs):
        for k in cls._methods:
#             if the method in cls._methods is in the `kwargs` of the instantiated class select the value of the kwargs (ie the new method)
            arg = kwargs.pop(k,None) #it returns None if it cannot find the key in the kwargs dict
            if arg is not None:
#                 do you want the func to be replaced to be a normal method and not a class method?
                if as_method: arg = method(arg)
#               if the method we want to replace with is a classmethod from another class, change to become 
#               a class method of the present class by binding it to the present class using `MethodType`
                if isinstance(arg,MethodType): arg = MethodType(arg.__func__, self)
#                 change the method in the class with this new method or classmethod
                setattr(self, k, arg)
#         set the signature if the present class init
        old_init(self, *args, **kwargs)
#     update the signature of the present init to prevent what happens with decorators and signatures
    functools.update_wrapper(_init, old_init)
#   set all the cls._methods to be kwargs with value of None so that we can pass all into signature
    cls.__init__ = use_kwargs(cls._methods)(_init)
    if hasattr(cls, '__signature__'): cls.__signature__ = _rm_self(inspect.signature(cls.__init__))
    return cls


#export
def funcs_kwargs(as_method=False):
    "Replace methods in `cls._methods` with those from `kwargs`"
    if callable(as_method): return _funcs_kwargs(as_method, False)
#     we have a closure so it is a decorator
    return partial(_funcs_kwargs, as_method=as_method)


# +
#tendo test

class New():
    def a(self):
        return 'I am the replacer from func a'
    
    @classmethod
    def b(self, a, **kwargs):
        return 'I am the replacer from func b'

@funcs_kwargs(as_method=True)
class _F():
    _methods=['bar']
    def __init__(self, **kwargs): assert not kwargs
    
    def foo(cls, a=1, b=2):
        return 'Hi'
    
#     @classmethod
#     def bar(cls, c=3, **kwargs):
#         return 'Hello'
        
def lambda_(a): return 'i am the replacer lambda_'

f = _F()
# print(f'Initial f.bar:    {f.bar().upper()}')
f = _F(bar = lambda_) #normal func replacement

test_eq(f.bar(), 'i am the replacer lambda_')


f = _F()
# print(f'Initial f.bar:    {f.bar().upper()}')
f = _F(bar = New.b) #classmethod replacement

test_eq(f.bar(), 'I am the replacer from func b')


# +
@funcs_kwargs
class T:
    _methods=['b']
    def __init__(self, f=1, **kwargs): assert not kwargs
    def a(self): return 1
    def b(self): return 2
    
t = T()
test_eq(t.a(), 1)
test_eq(t.b(), 2)
t = T(b = lambda:3)
test_eq(t.b(), 3)
# so long as the method is in `_methods`, it'll have a default value of None when you check the class's signature thus overriding the class' __init__ args
test_sig(T, '(f=1, *, b=None)') #the b=None happens due to the `use_kwargs` that was called in `_funcs_kwargs`
test_fail(lambda: T(a = lambda:3)) # a is not in T._methods

def _f(self,a=1): return a+1

@funcs_kwargs(True)
class T: _methods=['b']
t = T(b = _f)
test_eq(t.b(2), 3)

class T2(T):
    def __init__(self,a):
        super().__init__(b = lambda self:3) #calling super on `T` should override the `b` which was added to the `T` class by `func_kwargs`
        self.a=a
t = T2(a=1)
test_eq(t.b(), 3)
test_sig(T2, '(a)') #We have `T2`'s normal signature which is expected

def _g(a=1): return a+1
class T3(T): b = staticmethod(_g) #funcs_kwargs can be used to override staticmethods also (static methods by default aren't bound to the class)
t = T3()
test_eq(t.b(2), 3)


# +
#hide
#test it works with PrePostInitMeta
class A(metaclass=PrePostInitMeta): pass

@funcs_kwargs
class B(A):
    _methods = ['m1']
    def __init__(self, **kwargs): pass
    
test_sig(B, '(*, m1=None)')


# -

# Runtime type checking is handy, so let's make it easy!

@contextmanager
def working_directory(path):
    "Change working directory to `path` and return to previous on exit."
    prev_cwd = Path.cwd()
    os.chdir(path)
    try: yield
    finally: os.chdir(prev_cwd)


# +
#def is_listy(x): return isinstance(x,(list,tuple,Generator))

# +
class T:
    _methods=['b']
    def __init__(self, f=1, **kwargs): assert not kwargs
    @classmethod
    def a(self, v=8): return 1
    def b(self): return 2
    
vars(T)['a'].__func__


# +
# this func ensures that a class and all it's public methods are properly documented
#export
def add_docs(cls, cls_doc=None, **docs):
    "Copy values from `docs` to `cls` docstrings, and confirm all public methods are documented"
    if cls_doc is not None: cls.__doc__ = cls_doc
#     the `**docs` arg contains key-value pairs for method_name:method
    for k,v in docs.items():
#         get the method in the class from it's name
        f = getattr(cls,k)
#     check if the method has a __func__ attr if not, set it 
        if hasattr(f,'__func__'): f = f.__func__ # required for class methods
        f.__doc__ = v
        
#     check for methods/callables in the class's namespace that aren't inner methods(ie don't start with '_') 
#     and do not have docs(ie no `__doc__` attr) and list them
    # List of public callables without docstring
    nodoc = [c for n,c in vars(cls).items() if callable(c)
             and not n.startswith('_') and c.__doc__ is None]
#     if there is a callable without docs, raise an exception showing that callable
    assert not nodoc, f"Missing docs: {nodoc}"
#     if the cls has a `__doc__` but is None or empty raise an exception showing that the class has no doc
    assert cls.__doc__ is not None, f"Missing class docs: {cls}"


# -

#export
def docs(cls):
    "Decorator version of `add_docs`, using `_docs` dict"
    add_docs(cls, **cls._docs)
    return cls


# +
class _T:
    def f(self): pass
    @classmethod
    def g(cls): pass
add_docs(_T, cls_doc="a", f="f", g="g") #specify the docs as kwargs

test_eq(_T.__doc__, "a")
test_eq(_T.f.__doc__, "f")
test_eq(_T.g.__doc__, "g")


# -

# Tendo test
@docs
class _T:
    _docs = {'cls_doc':'a', 'f':'f', 'g':'g'} #specify the doc as a dict
    def f(self): pass
    @classmethod
    def g(cls): pass
# everything now has a doc
test_eq(_T.__doc__, "a")
test_eq(_T.f.__doc__, "f")
test_eq(_T.g.__doc__, "g")

show_doc(is_iter) # `show_doc` is a n nbdev method

assert is_iter([1])
assert not is_iter(array(1))
assert is_iter(array([1,2]))
assert (o for o in range(3))


#export
class _Arg:
    def __init__(self,i): self.i = i
arg0 = _Arg(0)
arg1 = _Arg(1)
arg2 = _Arg(2)
arg3 = _Arg(3)
arg4 = _Arg(4)


# +
# partial??
# -

#export
class bind:
    "Same as `partial`, except you can use `arg0` `arg1` etc param placeholders"
    def __init__(self, fn, *pargs, **pkwargs):
        self.fn,self.pargs,self.pkwargs = fn,pargs,pkwargs
        self.maxi = max((x.i for x in pargs if isinstance(x, _Arg)), default=-1)

    def __call__(self, *args, **kwargs):
        args = list(args)
        kwargs = {**self.pkwargs,**kwargs} #what is the purpose of kwargs tendo asks...doesnt't seem to do anything
        for k,v in kwargs.items():
            if isinstance(v,_Arg): kwargs[k] = args.pop(v.i)
        fargs = [args[x.i] if isinstance(x, _Arg) else x for x in self.pargs] + args[self.maxi+1:]
        return self.fn(*fargs, **kwargs)


def myfn(a,b,c,d=1,e=2): return(a,b,c,d,e)
test_eq(bind(myfn, arg1, 17, arg0, e=3)(19,14), (14,17,19,1,3))
test_eq(bind(myfn, 17, arg0, e=3)(19,14), (17,19,14,1,3))
test_eq(bind(myfn, 17, e=3)(19,14), (17,19,14,1,3))
test_eq(bind(myfn)(17,19,14), (17,19,14,1,2))
test_eq(bind(myfn, 17,19,14,e=arg0)(3), (17,19,14,1,3))


# ## GetAttr -

#export
def custom_dir(c, add:list):
    "Implement custom `__dir__`, adding `add` to `cls`"
    return dir(type(c)) + list(c.__dict__.keys()) + add


#export
class GetAttr:
    "Inherit from this to have all attr accesses in `self._xtra` passed down to `self.default`"
    _default='default'
    def _component_attr_filter(self,k):
        if k.startswith('__') or k in ('_xtra',self._default): return False
        xtra = getattr(self,'_xtra',None)
        return xtra is None or k in xtra
    def _dir(self): return [k for k in dir(getattr(self,self._default)) if self._component_attr_filter(k)]
    def __getattr__(self,k):
        if self._component_attr_filter(k):
            attr = getattr(self,self._default,None)
            if attr is not None: return getattr(attr,k)
        raise AttributeError(k)
    def __dir__(self): return custom_dir(self,self._dir())
#     def __getstate__(self): return self.__dict__
    def __setstate__(self,data): self.__dict__.update(data)


# Inherit from `GetAttr` to have attr access passed down to an instance attribute. 
# This makes it easy to create composites that don't require callers to know about their components.
#
# You can customise the behaviour of `GetAttr` in subclasses via;
# - `_default`
#     - By default, this is set to `'default'`, so attr access is passed down to `self.default`
#     - `_default` can be set to the name of any instance attribute that does not start with dunder `__`
# - `_xtra`
#     - By default, this is `None`, so all attr access is passed down
#     - You can limit which attrs get passed down by setting `_xtra` to a list of attribute names

# +
class _C(GetAttr):
    # allow all attributes to get passed to `self.default` (by leaving _xtra=None)
    def __init__(self,a): self.default = a
    def foo(self): noop

t = _C('Hi')
test_eq(t.lower(), 'hi')
test_eq(t.upper(), 'HI')
assert 'lower' in dir(t)
assert 'upper' in dir(t)


# +
class _C(GetAttr):
    _xtra = ['lower'] # specify which attributes get passed to `self.default`
    def __init__(self,a): self.default = a
    def foo(self): noop

t = _C('Hi')
test_eq(t.default, 'Hi')
test_eq(t.lower(), 'hi')
test_fail(lambda: t.upper())
assert 'lower' in dir(t)
assert 'upper' not in dir(t)


# +
class _C(GetAttr):
    _default = '_data' # use different component name; `self._data` rather than `self.default`
    def __init__(self,a): self._data = a
    def foo(self): noop

t = _C('Hi')
test_eq(t._data, 'Hi')
test_eq(t.lower(), 'hi')
test_eq(t.upper(), 'HI')
assert 'lower' in dir(t)
assert 'upper' in dir(t)


# -

class _C(GetAttr):
    _default = 'data' # use a bad component name; i.e. self.data does not exist
    def __init__(self,a): self.default = a
    def foo(self): noop
# TODO: should we raise an error when we create a new instance ...
t = _C('Hi')
test_eq(t.default, 'Hi')
# ... or is it enough for all GetAttr features to raise errors
test_fail(lambda: t.data)
test_fail(lambda: t.lower())
test_fail(lambda: t.upper())
test_fail(lambda: dir(t))


# +
#hide
# I don't think this test is essential to the docs but it probably makes sense to
# check that everything works when we set both _xtra and _default to non-default values
class _C(GetAttr):
    _xtra = ['lower', 'upper']
    _default = 'data'
    def __init__(self,a): self.data = a
    def foo(self): noop

t = _C('Hi')
test_eq(t.data, 'Hi')
test_eq(t.lower(), 'hi')
test_eq(t.upper(), 'HI')
assert 'lower' in dir(t)
assert 'upper' in dir(t)


# +
#hide
#  when consolidating the filter logic, I choose the previous logic from 
# __getattr__  k.startswith('__') rather than
# _dir         k.startswith('_'). 
class _C(GetAttr):
    def __init__(self): self.default = type('_D', (), {'_under': 1, '__dunder': 2})() 
    
t = _C()
test_eq(t.default._under, 1)
test_eq(t._under, 1)           # _ prefix attr access is allowed on component
assert '_under' in dir(t)

test_eq(t.default.__dunder, 2)
test_fail(lambda: t.__dunder)  # __ prefix attr access is not allowed on component
assert '__dunder' not in dir(t)

assert t.__dir__ is not None   # __ prefix attr access is allowed on composite
assert '__dir__' in dir(t)


# -

class B:
    def __init__(self): self.a = A()


@funcs_kwargs
class A(GetAttr):
    wif=after_iter= noops
    _methods = 'wif after_iter'.split()
    _default = 'dataset'
    def __init__(self, **kwargs): pass


a = A()
b = A(wif=a.wif)


# +
#Failing test. TODO Jeremy, not sure what you were testing here
#a = A()
#b = A(wif=a.wif)
#tst = pickle.dumps(b)
#c = pickle.loads(tst)
# -

#export
def delegate_attr(self, k, to):
    "Use in `__getattr__` to delegate to attr `to` without inheriting from `GetAttr`"
    if k.startswith('_') or k==to: raise AttributeError(k)
    try: return getattr(getattr(self,to), k)
    except AttributeError: raise AttributeError(k) from None


# +
class _C:
    f = 'Hi'
    def __getattr__(self, k): return delegate_attr(self, k, 'f')

t = _C()
test_eq(t.lower(), 'hi')


# -

# ## L -

# +
#export
def _is_array(x): return hasattr(x,'__array__') or hasattr(x,'iloc')

def _listify(o):
    if o is None: return []
    if isinstance(o, list): return o
    if isinstance(o, str) or _is_array(o): return [o]
    if is_iter(o): return list(o)
    return [o]


# -

# export
def coll_repr(c, max_n=10):
    "String repr of up to `max_n` items of (possibly lazy) collection `c`"
    return f'(#{len(c)}) [' + ','.join(itertools.islice(map(repr,c), max_n)) + (
        '...' if len(c)>10 else '') + ']'


test_eq(coll_repr(range(1000), 5), '(#1000) [0,1,2,3,4...]')


# export
def mask2idxs(mask):
    "Convert bool mask or index list to index `L`"
    if isinstance(mask,slice): return mask
    mask = list(mask)
    if len(mask)==0: return []
    it = mask[0]
    if hasattr(it,'item'): it = it.item()
    if isinstance(it,(bool,NoneType,np.bool_)): return [i for i,m in enumerate(mask) if m]
    return [int(i) for i in mask]


# just for tests
import torch

test_eq(mask2idxs([False,True,False,True]), [1,3])
test_eq(mask2idxs(array([False,True,False,True])), [1,3])
test_eq(mask2idxs(torch.tensor([False,True,False,True])), [1,3])
test_eq(mask2idxs(array([1,2,3])), [1,2,3])

#export
listable_types = typing.Collection,Generator,map,filter,zip


#export
class CollBase:
    "Base class for composing a list of `items`"
    def __init__(self, items): self.items = items
    def __len__(self): return len(self.items)
    def __getitem__(self, k): return self.items[list(k) if isinstance(k,CollBase) else k]
    def __setitem__(self, k, v): self.items[list(k) if isinstance(k,CollBase) else k] = v
    def __delitem__(self, i): del(self.items[i])
    def __repr__(self): return self.items.__repr__()
    def __iter__(self): return self.items.__iter__()


#export
def cycle(o):
    "Like `itertools.cycle` except creates list of `None`s if `o` is empty"
    o = _listify(o)
    return itertools.cycle(o) if o is not None and len(o) > 0 else itertools.cycle([None])


test_eq(itertools.islice(cycle([1,2,3]),5), [1,2,3,1,2])
test_eq(itertools.islice(cycle([]),3), [None]*3)
test_eq(itertools.islice(cycle(None),3), [None]*3)
test_eq(itertools.islice(cycle(1),3), [1,1,1])


#export
def zip_cycle(x, *args):
    "Like `itertools.zip_longest` but `cycle`s through elements of all but first argument"
    return zip(x, *map(cycle,args))


test_eq(zip_cycle([1,2,3,4],list('abc')), [(1, 'a'), (2, 'b'), (3, 'c'), (4, 'a')])


#export
def is_indexer(idx):
    "Test whether `idx` will index a single item in a list"
    return isinstance(idx,int) or not getattr(idx,'ndim',1)


#export
def negate_func(f):
    "Create new function that negates result of `f`"
    def _f(*args, **kwargs): return not f(*args, **kwargs)
    return _f


def f(a): return a>0
test_eq(f(1),True)
test_eq(negate_func(f)(1),False)
test_eq(negate_func(f)(a=-1),True)


#export
class L(CollBase):
    "Behaves like a list of `items` but can also index with list of indices or masks"
    _default='items'
    def __init__(self, items=None, *rest, use_list=False, match=None):
        if rest: items = (items,)+rest
        if items is None: items = []
        if (use_list is not None) or not _is_array(items):
            items = list(items) if use_list else _listify(items)
        if match is not None:
            if is_coll(match): match = len(match)
            if len(items)==1: items = items*match
            else: assert len(items)==match, 'Match length mismatch'
        super().__init__(items)

    @property
    def _xtra(self): return None
    def _new(self, items, *args, **kwargs): return type(self)(items, *args, use_list=None, **kwargs)
    def __getitem__(self, idx): return self._get(idx) if is_indexer(idx) else L(self._get(idx), use_list=None)
    def copy(self): return self._new(self.items.copy())

    def _get(self, i):
        if is_indexer(i) or isinstance(i,slice): return getattr(self.items,'iloc',self.items)[i]
        i = mask2idxs(i)
        return (self.items.iloc[list(i)] if hasattr(self.items,'iloc')
                else self.items.__array__()[(i,)] if hasattr(self.items,'__array__')
                else [self.items[i_] for i_ in i])

    def __setitem__(self, idx, o):
        "Set `idx` (can be list of indices, or mask, or int) items to `o` (which is broadcast if not iterable)"
        if isinstance(idx, int): self.items[idx] = o
        else:
            idx = idx if isinstance(idx,L) else _listify(idx)
            if not is_iter(o): o = [o]*len(idx)
            for i,o_ in zip(idx,o): self.items[i] = o_

    def __iter__(self): return iter(self.items.itertuples() if hasattr(self.items,'iloc') else self.items)
    def __contains__(self,b): return b in self.items
    def __reversed__(self): return self._new(reversed(self.items))
    def __invert__(self): return self._new(not i for i in self)
    def __eq__(self,b): return False if isinstance(b, (str,dict,set)) else all_equal(b,self)
    def __repr__(self): return repr(self.items) if _is_array(self.items) else coll_repr(self)
    def __mul__ (a,b): return a._new(a.items*b)
    def __add__ (a,b): return a._new(a.items+_listify(b))
    def __radd__(a,b): return a._new(b)+a
    def __addi__(a,b):
        a.items += list(b)
        return a

    def sorted(self, key=None, reverse=False):
        if isinstance(key,str):   k=lambda o:getattr(o,key,0)
        elif isinstance(key,int): k=itemgetter(key)
        else: k=key
        return self._new(sorted(self.items, key=k, reverse=reverse))

    @classmethod
    def split(cls, s, sep=None, maxsplit=-1): return cls(s.split(sep,maxsplit))

    @classmethod
    def range(cls, a, b=None, step=None):
        if is_coll(a): a = len(a)
        return cls(range(a,b,step) if step is not None else range(a,b) if b is not None else range(a))

    def map(self, f, *args, **kwargs):
        g = (bind(f,*args,**kwargs) if callable(f)
             else f.format if isinstance(f,str)
             else f.__getitem__)
        return self._new(map(g, self))

    def filter(self, f, negate=False, **kwargs):
        if kwargs: f = partial(f,**kwargs)
        if negate: f = negate_func(f)
        return self._new(filter(f, self))

    def argwhere(self, f, negate=False, **kwargs):
        if kwargs: f = partial(f,**kwargs)
        if negate: f = negate_func(f)
        return self._new(i for i,o in enumerate(self) if f(o))

    def unique(self): return L(dict.fromkeys(self).keys())
    def enumerate(self): return L(enumerate(self))
    def val2idx(self): return {v:k for k,v in self.enumerate()}
    def itemgot(self, *idxs):
        x = self
        for idx in idxs: x = x.map(itemgetter(idx))
        return x

    def attrgot(self, k, default=None): return self.map(lambda o:getattr(o,k,default))
    def cycle(self): return cycle(self)
    def map_dict(self, f=noop, *args, **kwargs): return {k:f(k, *args,**kwargs) for k in self}
    def starmap(self, f, *args, **kwargs): return self._new(itertools.starmap(partial(f,*args,**kwargs), self))
    def zip(self, cycled=False): return self._new((zip_cycle if cycled else zip)(*self))
    def zipwith(self, *rest, cycled=False): return self._new([self, *rest]).zip(cycled=cycled)
    def map_zip(self, f, *args, cycled=False, **kwargs): return self.zip(cycled=cycled).starmap(f, *args, **kwargs)
    def map_zipwith(self, f, *rest, cycled=False, **kwargs): return self.zipwith(*rest, cycled=cycled).starmap(f, **kwargs)
    def concat(self): return self._new(itertools.chain.from_iterable(self.map(L)))
    def shuffle(self):
        it = copy(self.items)
        random.shuffle(it)
        return self._new(it)

    def append(self,o): return self.items.append(o)
    def remove(self,o): return self.items.remove(o)
    def count (self,o): return self.items.count(o)
    def reverse(self ): return self.items.reverse()
    def pop(self,o=-1): return self.items.pop(o)
    def clear(self   ): return self.items.clear()
    def index(self, value, start=0, stop=sys.maxsize): return self.items.index(value, start, stop)
    def sort(self, key=None, reverse=False): return self.items.sort(key=key, reverse=reverse)
    def reduce(self, f, initial=None): return reduce(f, self) if initial is None else reduce(f, self, initial)
    def sum(self): return self.reduce(operator.add)
    def product(self): return self.reduce(operator.mul)


inspect.signature(L)

#export
_docs = {o:"Passthru to `list` method" for o in
         'append count remove reverse sort pop clear index'.split()}
add_docs(L,
         __getitem__="Retrieve `idx` (can be list of indices, or mask, or int) items",
         range="Same as `range`, but returns an `L`. Can pass a collection for `a`, to use `len(a)`",
         split="Same as `str.split`, but returns an `L`",
         copy="Same as `list.copy`, but returns an `L`",
         sorted="New `L` sorted by `key`. If key is str then use `attrgetter`. If key is int then use `itemgetter`",
         unique="Unique items, in stable order",
         val2idx="Dict from value to index",
         filter="Create new `L` filtered by predicate `f`, passing `args` and `kwargs` to `f`",
         argwhere="Like `filter`, but return indices for matching items",
         map="Create new `L` with `f` applied to all `items`, passing `args` and `kwargs` to `f`",
         map_dict="Like `map`, but creates a dict from `items` to function results",
         starmap="Like `map`, but use `itertools.starmap`",
         itemgot="Create new `L` with item `idx` of all `items`",
         attrgot="Create new `L` with attr `k` of all `items`",
         cycle="Same as `itertools.cycle`",
         enumerate="Same as `enumerate`",
         zip="Create new `L` with `zip(*items)`",
         zipwith="Create new `L` with `self` zip with each of `*rest`",
         map_zip="Combine `zip` and `starmap`",
         map_zipwith="Combine `zipwith` and `starmap`",
         concat="Concatenate all elements of list",
         shuffle="Same as `random.shuffle`, but not inplace",
         reduce="Wrapper for `functools.reduce`",
         sum="Sum of the items",
         product="Product of the items",
         **_docs)

#export
Sequence.register(L);

# You can create an `L` from an existing iterable (e.g. a list, range, etc) and access or modify it with an int list/tuple index, mask, int, or slice. All `list` methods can also be used with `L`.

t = L(range(12))
test_eq(t, list(range(12)))
test_ne(t, list(range(11)))
t.reverse()
test_eq(t[0], 11)
t[3] = "h"
test_eq(t[3], "h")
t[3,5] = ("j","k")
test_eq(t[3,5], ["j","k"])
test_eq(t, L(t))
test_eq(L(L(1,2),[3,4]), ([1,2],[3,4]))
t



StarcoderdataJupyterScriptsDedupFiltered