@@ -777,6 +777,236 @@ def test_local_kinds(self):
777
777
kinds = _testinternalcapi .get_co_localskinds (func .__code__ )
778
778
self .assertEqual (kinds , expected )
779
779
780
+ @unittest .skipIf (_testinternalcapi is None , "missing _testinternalcapi" )
781
+ def test_var_counts (self ):
782
+ self .maxDiff = None
783
+ def new_var_counts (* ,
784
+ posonly = 0 ,
785
+ posorkw = 0 ,
786
+ kwonly = 0 ,
787
+ varargs = 0 ,
788
+ varkwargs = 0 ,
789
+ purelocals = 0 ,
790
+ argcells = 0 ,
791
+ othercells = 0 ,
792
+ freevars = 0 ,
793
+ globalvars = 0 ,
794
+ attrs = 0 ,
795
+ unknown = 0 ,
796
+ ):
797
+ nargvars = posonly + posorkw + kwonly + varargs + varkwargs
798
+ nlocals = nargvars + purelocals + othercells
799
+ if isinstance (globalvars , int ):
800
+ globalvars = {
801
+ 'total' : globalvars ,
802
+ 'numglobal' : 0 ,
803
+ 'numbuiltin' : 0 ,
804
+ 'numunknown' : globalvars ,
805
+ }
806
+ else :
807
+ g_numunknown = 0
808
+ if isinstance (globalvars , dict ):
809
+ numglobal = globalvars ['numglobal' ]
810
+ numbuiltin = globalvars ['numbuiltin' ]
811
+ size = 2
812
+ if 'numunknown' in globalvars :
813
+ g_numunknown = globalvars ['numunknown' ]
814
+ size += 1
815
+ assert len (globalvars ) == size , globalvars
816
+ else :
817
+ assert not isinstance (globalvars , str ), repr (globalvars )
818
+ try :
819
+ numglobal , numbuiltin = globalvars
820
+ except ValueError :
821
+ numglobal , numbuiltin , g_numunknown = globalvars
822
+ globalvars = {
823
+ 'total' : numglobal + numbuiltin + g_numunknown ,
824
+ 'numglobal' : numglobal ,
825
+ 'numbuiltin' : numbuiltin ,
826
+ 'numunknown' : g_numunknown ,
827
+ }
828
+ unbound = globalvars ['total' ] + attrs + unknown
829
+ return {
830
+ 'total' : nlocals + freevars + unbound ,
831
+ 'locals' : {
832
+ 'total' : nlocals ,
833
+ 'args' : {
834
+ 'total' : nargvars ,
835
+ 'numposonly' : posonly ,
836
+ 'numposorkw' : posorkw ,
837
+ 'numkwonly' : kwonly ,
838
+ 'varargs' : varargs ,
839
+ 'varkwargs' : varkwargs ,
840
+ },
841
+ 'numpure' : purelocals ,
842
+ 'cells' : {
843
+ 'total' : argcells + othercells ,
844
+ 'numargs' : argcells ,
845
+ 'numothers' : othercells ,
846
+ },
847
+ 'hidden' : {
848
+ 'total' : 0 ,
849
+ 'numpure' : 0 ,
850
+ 'numcells' : 0 ,
851
+ },
852
+ },
853
+ 'numfree' : freevars ,
854
+ 'unbound' : {
855
+ 'total' : unbound ,
856
+ 'globals' : globalvars ,
857
+ 'numattrs' : attrs ,
858
+ 'numunknown' : unknown ,
859
+ },
860
+ }
861
+
862
+ import test ._code_definitions as defs
863
+ funcs = {
864
+ defs .spam_minimal : new_var_counts (),
865
+ defs .spam_full : new_var_counts (
866
+ posonly = 2 ,
867
+ posorkw = 2 ,
868
+ kwonly = 2 ,
869
+ varargs = 1 ,
870
+ varkwargs = 1 ,
871
+ purelocals = 4 ,
872
+ globalvars = 3 ,
873
+ attrs = 1 ,
874
+ ),
875
+ defs .spam : new_var_counts (
876
+ posorkw = 1 ,
877
+ ),
878
+ defs .spam_N : new_var_counts (
879
+ posorkw = 1 ,
880
+ purelocals = 1 ,
881
+ ),
882
+ defs .spam_C : new_var_counts (
883
+ posorkw = 1 ,
884
+ purelocals = 1 ,
885
+ argcells = 1 ,
886
+ othercells = 1 ,
887
+ ),
888
+ defs .spam_NN : new_var_counts (
889
+ posorkw = 1 ,
890
+ purelocals = 1 ,
891
+ ),
892
+ defs .spam_NC : new_var_counts (
893
+ posorkw = 1 ,
894
+ purelocals = 1 ,
895
+ argcells = 1 ,
896
+ othercells = 1 ,
897
+ ),
898
+ defs .spam_CN : new_var_counts (
899
+ posorkw = 1 ,
900
+ purelocals = 1 ,
901
+ argcells = 1 ,
902
+ othercells = 1 ,
903
+ ),
904
+ defs .spam_CC : new_var_counts (
905
+ posorkw = 1 ,
906
+ purelocals = 1 ,
907
+ argcells = 1 ,
908
+ othercells = 1 ,
909
+ ),
910
+ defs .eggs_nested : new_var_counts (
911
+ posorkw = 1 ,
912
+ ),
913
+ defs .eggs_closure : new_var_counts (
914
+ posorkw = 1 ,
915
+ freevars = 2 ,
916
+ ),
917
+ defs .eggs_nested_N : new_var_counts (
918
+ posorkw = 1 ,
919
+ purelocals = 1 ,
920
+ ),
921
+ defs .eggs_nested_C : new_var_counts (
922
+ posorkw = 1 ,
923
+ purelocals = 1 ,
924
+ argcells = 1 ,
925
+ freevars = 2 ,
926
+ ),
927
+ defs .eggs_closure_N : new_var_counts (
928
+ posorkw = 1 ,
929
+ purelocals = 1 ,
930
+ freevars = 2 ,
931
+ ),
932
+ defs .eggs_closure_C : new_var_counts (
933
+ posorkw = 1 ,
934
+ purelocals = 1 ,
935
+ argcells = 1 ,
936
+ othercells = 1 ,
937
+ freevars = 2 ,
938
+ ),
939
+ defs .ham_nested : new_var_counts (
940
+ posorkw = 1 ,
941
+ ),
942
+ defs .ham_closure : new_var_counts (
943
+ posorkw = 1 ,
944
+ freevars = 3 ,
945
+ ),
946
+ defs .ham_C_nested : new_var_counts (
947
+ posorkw = 1 ,
948
+ ),
949
+ defs .ham_C_closure : new_var_counts (
950
+ posorkw = 1 ,
951
+ freevars = 4 ,
952
+ ),
953
+ }
954
+ assert len (funcs ) == len (defs .FUNCTIONS ), (len (funcs ), len (defs .FUNCTIONS ))
955
+ for func in defs .FUNCTIONS :
956
+ with self .subTest (func ):
957
+ expected = funcs [func ]
958
+ counts = _testinternalcapi .get_code_var_counts (func .__code__ )
959
+ self .assertEqual (counts , expected )
960
+
961
+ def func_with_globals_and_builtins ():
962
+ mod1 = _testinternalcapi
963
+ mod2 = dis
964
+ mods = (mod1 , mod2 )
965
+ checks = tuple (callable (m ) for m in mods )
966
+ return callable (mod2 ), tuple (mods ), list (mods ), checks
967
+
968
+ func = func_with_globals_and_builtins
969
+ with self .subTest (f'{ func } code' ):
970
+ expected = new_var_counts (
971
+ purelocals = 4 ,
972
+ globalvars = 5 ,
973
+ )
974
+ counts = _testinternalcapi .get_code_var_counts (func .__code__ )
975
+ self .assertEqual (counts , expected )
976
+
977
+ with self .subTest (f'{ func } with own globals and builtins' ):
978
+ expected = new_var_counts (
979
+ purelocals = 4 ,
980
+ globalvars = (2 , 3 ),
981
+ )
982
+ counts = _testinternalcapi .get_code_var_counts (func )
983
+ self .assertEqual (counts , expected )
984
+
985
+ with self .subTest (f'{ func } without globals' ):
986
+ expected = new_var_counts (
987
+ purelocals = 4 ,
988
+ globalvars = (0 , 3 , 2 ),
989
+ )
990
+ counts = _testinternalcapi .get_code_var_counts (func , globalsns = {})
991
+ self .assertEqual (counts , expected )
992
+
993
+ with self .subTest (f'{ func } without both' ):
994
+ expected = new_var_counts (
995
+ purelocals = 4 ,
996
+ globalvars = 5 ,
997
+ )
998
+ counts = _testinternalcapi .get_code_var_counts (func , globalsns = {},
999
+ builtinsns = {})
1000
+ self .assertEqual (counts , expected )
1001
+
1002
+ with self .subTest (f'{ func } without builtins' ):
1003
+ expected = new_var_counts (
1004
+ purelocals = 4 ,
1005
+ globalvars = (2 , 0 , 3 ),
1006
+ )
1007
+ counts = _testinternalcapi .get_code_var_counts (func , builtinsns = {})
1008
+ self .assertEqual (counts , expected )
1009
+
780
1010
781
1011
def isinterned (s ):
782
1012
return s is sys .intern (('_' + s + '_' )[1 :- 1 ])
0 commit comments