File size: 4,106 Bytes
b7731cd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#-----------------------------------------------------------------------------
#  Copyright (C) 2013 Min RK
#
#  Distributed under the terms of the 2-clause BSD License.
#-----------------------------------------------------------------------------

from contextlib import contextmanager

import ctypes
import ctypes.util

objc = ctypes.cdll.LoadLibrary(ctypes.util.find_library('objc'))
_ = ctypes.cdll.LoadLibrary(ctypes.util.find_library('Foundation'))

void_p = ctypes.c_void_p
ull = ctypes.c_uint64

objc.objc_getClass.restype = void_p
objc.sel_registerName.restype = void_p
objc.objc_msgSend.restype = void_p
objc.objc_msgSend.argtypes = [void_p, void_p]

msg = objc.objc_msgSend

def _utf8(s):
    """ensure utf8 bytes"""
    if not isinstance(s, bytes):
        s = s.encode('utf8')
    return s

def n(name):
    """create a selector name (for methods)"""
    return objc.sel_registerName(_utf8(name))

def C(classname):
    """get an ObjC Class by name"""
    ret = objc.objc_getClass(_utf8(classname))
    assert ret is not None, "Couldn't find Class %s" % classname
    return ret

# constants from Foundation

NSActivityIdleDisplaySleepDisabled             = (1 << 40)
NSActivityIdleSystemSleepDisabled              = (1 << 20)
NSActivitySuddenTerminationDisabled            = (1 << 14)
NSActivityAutomaticTerminationDisabled         = (1 << 15)
NSActivityUserInitiated                        = (0x00FFFFFF | NSActivityIdleSystemSleepDisabled)
NSActivityUserInitiatedAllowingIdleSystemSleep = (NSActivityUserInitiated & ~NSActivityIdleSystemSleepDisabled)
NSActivityBackground                           = 0x000000FF
NSActivityLatencyCritical                      = 0xFF00000000

def beginActivityWithOptions(options, reason=""):
    """Wrapper for:
    
    [ [ NSProcessInfo processInfo] 
        beginActivityWithOptions: (uint64)options
                          reason: (str)reason
    ]
    """
    NSProcessInfo = C('NSProcessInfo')
    NSString = C('NSString')
    
    objc.objc_msgSend.argtypes = [void_p, void_p, void_p]
    reason = msg(NSString, n("stringWithUTF8String:"), _utf8(reason))
    objc.objc_msgSend.argtypes = [void_p, void_p]
    info = msg(NSProcessInfo, n('processInfo'))
    objc.objc_msgSend.argtypes = [void_p, void_p, ull, void_p]
    activity = msg(info,
        n('beginActivityWithOptions:reason:'),
        ull(options),
        void_p(reason)
    )
    return activity

def endActivity(activity):
    """end a process activity assertion"""
    NSProcessInfo = C('NSProcessInfo')
    objc.objc_msgSend.argtypes = [void_p, void_p]
    info = msg(NSProcessInfo, n('processInfo'))
    objc.objc_msgSend.argtypes = [void_p, void_p, void_p]
    msg(info, n("endActivity:"), void_p(activity))

_theactivity = None

def nope():
    """disable App Nap by setting NSActivityUserInitiatedAllowingIdleSystemSleep"""
    global _theactivity
    _theactivity = beginActivityWithOptions(
        NSActivityUserInitiatedAllowingIdleSystemSleep,
        "Because Reasons"
    )

def nap():
    """end the caffeinated state started by `nope`"""
    global _theactivity
    if _theactivity is not None:
        endActivity(_theactivity)
        _theactivity = None

def napping_allowed():
    """is napping allowed?"""
    return _theactivity is None

@contextmanager
def nope_scope(
        options=NSActivityUserInitiatedAllowingIdleSystemSleep,
        reason="Because Reasons"
    ):
    """context manager for beginActivityWithOptions.
    
    Within this context, App Nap will be disabled.
    """
    activity = beginActivityWithOptions(options, reason)
    try:
        yield
    finally:
        endActivity(activity)

__all__ = [
    "NSActivityIdleDisplaySleepDisabled",
    "NSActivityIdleSystemSleepDisabled",
    "NSActivitySuddenTerminationDisabled",
    "NSActivityAutomaticTerminationDisabled",
    "NSActivityUserInitiated",
    "NSActivityUserInitiatedAllowingIdleSystemSleep",
    "NSActivityBackground",
    "NSActivityLatencyCritical",
    "beginActivityWithOptions",
    "endActivity",
    "nope",
    "nap",
    "napping_allowed",
    "nope_scope",
]