44import  re 
55import  unicodedata 
66
7+ from  typing  import  Optional , Union , List , Any , Dict 
78from  lighttree .tree  import  Tree 
89
910
10- def  is_valid_attr_name (item ):
11+ def  is_valid_attr_name (item )  ->   bool :
1112    if  not  isinstance (item , str ):
1213        return  False 
1314    if  item .startswith ("__" ):
@@ -19,7 +20,7 @@ def is_valid_attr_name(item):
1920    return  True 
2021
2122
22- def  _coerce_attr (attr ):
23+ def  _coerce_attr (attr )  ->   Union [ str ,  None ] :
2324    if  not  len (attr ):
2425        return  None 
2526    new_attr  =  unicodedata .normalize ("NFD" , attr ).encode ("ASCII" , "ignore" ).decode ()
@@ -54,9 +55,9 @@ class Obj(object):
5455    _STRING_KEY_CONSTRAINT  =  True 
5556    _COERCE_ATTR  =  False 
5657
57-     def  __init__ (self , ** kwargs ):
58+     def  __init__ (self , ** kwargs )  ->   None :
5859        # will store non-valid names 
59-         self .__d  =  dict ()
60+         self .__d :  Dict [ str ,  Any ]  =  dict ()
6061        for  k , v  in  kwargs .items ():
6162            if  not  (
6263                isinstance (k , str )
@@ -95,21 +96,21 @@ def __setitem__(self, key, value):
9596        else :
9697            super (Obj , self ).__setattr__ (key , value )
9798
98-     def  __keys (self ):
99+     def  __keys (self )  ->   List [ Any ] :
99100        return  list (self .__d .keys ()) +  [
100101            k  for  k  in  self .__dict__ .keys () if  k  not  in "_REPR_NAME" , "_Obj__d" )
101102        ]
102103
103-     def  __contains__ (self , item ):
104+     def  __contains__ (self , item )  ->   bool :
104105        return  item  in  self .__keys ()
105106
106-     def  __str__ (self ):
107+     def  __str__ (self )  ->   str :
107108        return  "<%s> %s"  %  (
108109            str (self .__class__ ._REPR_NAME  or  self .__class__ .__name__ ),
109110            str (sorted (map (str , self .__keys ()))),
110111        )
111112
112-     def  __repr__ (self ):
113+     def  __repr__ (self )  ->   str :
113114        return  self .__str__ ()
114115
115116
@@ -124,18 +125,20 @@ class TreeBasedObj(Obj):
124125    _COERCE_ATTR  =  False 
125126    _ATTR  =  None 
126127
127-     def  __init__ (self , tree , root_path = None , depth = 1 , initial_tree = None ):
128+     def  __init__ (
129+         self ,
130+         tree : Tree ,
131+         root_path : Optional [str ] =  None ,
132+         depth : int  =  1 ,
133+         initial_tree : Optional [Tree ] =  None ,
134+     ) ->  None :
128135        super (TreeBasedObj , self ).__init__ ()
129-         if  not  isinstance (tree , Tree ):
130-             raise  ValueError (
131-                 'tree must be an instance of "lighttree.tree.Tree", got %s'  %  type (tree )
132-             )
133136        self ._tree  =  tree 
134137        self ._root_path  =  root_path 
135138        self ._initial_tree  =  initial_tree  if  initial_tree  is  not None  else  tree 
136139        self ._expand_attrs (depth )
137140
138-     def  _clone (self , nid , root_path , depth ) :
141+     def  _clone (self , nid :  str , root_path :  str , depth :  int )  ->   "TreeBasedObj" :
139142        _ , st  =  self ._tree .subtree (nid )
140143        return  self .__class__ (
141144            tree = st ,
@@ -144,24 +147,29 @@ def _clone(self, nid, root_path, depth):
144147            initial_tree = self ._initial_tree ,
145148        )
146149
147-     def  _expand_attrs (self , depth ) :
150+     def  _expand_attrs (self , depth :  int )  ->   None :
148151        if  depth :
149-             for  child_key , child_node  in  self ._tree .children (nid = self ._tree .root ):
152+             r  =  self ._tree .root 
153+             if  r  is  None :
154+                 return 
155+             for  child_key , child_node  in  self ._tree .children (nid = r ):
150156                if  self ._ATTR :
151157                    child_key  =  getattr (child_node , self ._ATTR )
158+                 str_child_key : str 
152159                if  isinstance (child_key , str ):
153160                    if  self ._COERCE_ATTR :
154161                        # if invalid coercion, coerce returns None, in this case we keep inital naming 
155162                        str_child_key  =  _coerce_attr (child_key ) or  child_key 
156163                    else :
157164                        str_child_key  =  child_key 
158165                else :
159-                     str_child_key  =  "i%d "  %  child_key 
166+                     str_child_key  =  "i%s "  %  child_key 
160167
168+                 child_root : str 
161169                if  self ._root_path  is  not None :
162170                    child_root  =  "%s.%s"  %  (self ._root_path , child_key )
163171                else :
164-                     child_root  =  child_key 
172+                     child_root  =  str ( child_key )  if   child_key   else   "" 
165173                self [str_child_key ] =  self ._clone (
166174                    child_node .identifier , root_path = child_root , depth = depth  -  1 
167175                )
@@ -182,7 +190,7 @@ def __getitem__(self, item):
182190            r ._expand_attrs (depth = 1 )
183191        return  r 
184192
185-     def  _show (self , * args , ** kwargs ):
193+     def  _show (self , * args , ** kwargs )  ->   str :
186194        tree_repr  =  self ._tree .show (* args , ** kwargs )
187195        if  self ._root_path  is  None :
188196            return  "<%s>\n %s"  %  (
@@ -196,8 +204,8 @@ def _show(self, *args, **kwargs):
196204            str (tree_repr ),
197205        )
198206
199-     def  __call__ (self , * args , ** kwargs ):
207+     def  __call__ (self , * args , ** kwargs )  ->   Tree :
200208        return  self ._tree 
201209
202-     def  __str__ (self ):
210+     def  __str__ (self )  ->   str :
203211        return  self ._show ()
0 commit comments